@mlx-bun/inference 0.0.0-stage → 1.0.0

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

@@ -0,0 +1,123 @@
+import { captureFullKvDonorRows } from "./full-kv-row-donor";
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import { appendFullKvRows } from "./full-kv-row-append";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { unrotateValues } from "../kernels/turboquant/ops";
+import { type AttentionRead, type Cache, type CommittedAttentionCache } from "../contracts/mlx/cache";
+import { kvStorage, type RotatedValueAttentionState } from "./kv-storage";
+import { TurboQuantKVCache } from "./turboquant-kv";
+import { BatchedTurboQuantKVCache } from "./batched-turboquant-kv";
+import { FullTransitioningKvRows } from "./full-transitioning-kv-rows";
+import { FullPrefillRow } from "./full-prefill-row";
+import { KVCache } from "./kv";
+import { Bf16FirstTurboQuantKVCache, completeConversion, convertToTurboQuant, transitionDue } from "./bf16-first-kv";
+import type { TurboQuantKernels } from "./turboquant-attention";
+
+/** The TurboQuant transition of `Bf16FirstTurboQuantKVCache` over
+ * full-attention rows: a bf16 row (`KVCache`, or a padded `FullPrefillRow`
+ * over one) encodes its live window once it holds `start` positions. */
+function settleTurboQuantRows(rows: Cache[], kBits: number, vBits: number, start: number, fusedDecode: boolean): void {
+  for (let index = 0; index < rows.length; index++) {
+    const row = rows[index]!;
+    if (row instanceof FullPrefillRow) {
+      const source = row.turboConversion;
+      if (!source || !transitionDue(source.offset, start)) continue;
+      const offset = source.offset;
+      rows[index] = completeConversion(source.toTurboQuantized(kBits, vBits, fusedDecode), offset, start);
+    } else if (row instanceof KVCache && transitionDue(row.offset, start)) {
+      rows[index] = convertToTurboQuant(row, kBits, vBits, start, fusedDecode);
+    }
+  }
+}
+
+/** The rows of the TurboQuant bf16-first lego (`Bf16FirstTurboQuantKVCache`)
+ * for continuous batching: TurboQuant tensor access and value-domain handling
+ * over the shared row-transition lifecycle (when a row converts is described on
+ * `TransitioningKvRows`). Model projections and attention remain batched.
+ * `fusedDecode` is resolved by whoever builds the layout. `turboQuantKernels` is
+ * the read of the graph whose caches these rows hold (dense for Qwen 3.5,
+ * deferred for Gemma 4 and MiniCPM5); the rows this layout makes and the
+ * layout it packs them into carry it, and a converted row attends with its
+ * source's.
+ *
+ * The reads are the graph's call today: this layout's mask before the append
+ * (no sliding window), then the kernel set's fetch and attention. The dense set
+ * fetches every row decoded (`updateAndFetch`) and runs the fused SDPA, Qwen
+ * 3.5's call; the deferred set fetches converted rows rotated
+ * (`updateAndFetchDeferredV`), runs the fused SDPA and applies the value
+ * transform that fetch captured, Gemma 4's and MiniCPM5's. */
+export class DelayedTurboQuantKVCache extends FullTransitioningKvRows<BatchedTurboQuantKVCache>
+  implements RotatedValueAttentionState, CommittedAttentionCache {
+  #rotated: boolean[] = [];
+  constructor(readonly kBits: number, readonly vBits: number, readonly start: number,
+    readonly fusedDecode: boolean, readonly turboQuantKernels: TurboQuantKernels, 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-turboquant:${kBits}:${vBits}:${start}`, conversionOffset: start,
+      settle: rows => settleTurboQuantRows(rows, kBits, vBits, start, fusedDecode),
+      converted: row => row instanceof TurboQuantKVCache,
+      makeLayout: () => new BatchedTurboQuantKVCache(kBits, vBits, turboQuantKernels, fusedDecode),
+      lone: storage => new Bf16FirstTurboQuantKVCache(kBits, vBits, start, fusedDecode, turboQuantKernels, undefined, storage),
+      ...(compose ? { compose } : {}) }, row);
+  }
+  /** Rows made for an empty layout carry the graph's read into their conversion. */
+  protected override plainRow(): KVCache { return new KVCache(undefined, this.turboQuantKernels); }
+  captureDonorRows(): import("../contracts/mlx/cache").KvDonorRows {
+    return this.packed?.captureDonorRows() ?? captureFullKvDonorRows(this.rows, this.leftPad, this.offset);
+  }
+  captureDonorAttention(): import("../contracts/mlx/cache").KvDonorAttention {
+    return decodedKvDonorAttention(this.captureDonorRows());
+  }
+  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v); }
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v); }
+  /** The committed read is the window read: a committed span reads this layout
+   * through the kernel set's fetch and the fused SDPA under the mask the layout
+   * builds, as a window does. */
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead { return this.#read(k, v); }
+  /** The read: the mask from this layout before the append (no sliding
+   * window), then the kernel set's fetch and attention. */
+  #read(k: MlxArray, v: MlxArray): AttentionRead {
+    const mask = this.makeMask(k.shape[2]!, null);
+    return this.turboQuantKernels.read({
+      decoded: () => this.updateAndFetch(k, v),
+      rotated: () => {
+        const [keys, values] = this.updateAndFetchDeferredV(k, v);
+        return { keys, values, restore: this.captureValueTransform() };
+      },
+    }, { mask, release: () => mask.arr?.dispose() });
+  }
+
+  #append(k: MlxArray, v: MlxArray, deferred: boolean): [MlxArray, MlxArray] {
+    this.advance();
+    if (this.packed) {
+      this.#rotated = Array.from({ length: this.batchSize! }, () => true);
+      return deferred ? this.packed.updateAndFetchDeferredV(k, v) : this.packed.updateAndFetch(k, v);
+    }
+    this.#rotated = this.rows.map(row => !!kvStorage(row).rotatedValueAttention);
+    const result = appendFullKvRows(this.rows, this.leftPad, this.offset + k.shape[2]!, k, v, deferred);
+    try { this.syncPositions(true); } catch (error) { result[0].dispose(); result[1].dispose(); throw error; }
+    return result;
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, false); }
+  updateAndFetchDeferredV(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#append(k, v, true); }
+  captureValueTransform(): (output: MlxArray) => MlxArray {
+    const rotated = [...this.#rotated];
+    if (rotated.every(value => value)) return unrotateValues;
+    if (rotated.every(value => !value)) return output => ops.contiguous(output);
+    return output => {
+      using restored = unrotateValues(output);
+      using flags = ops.fromInt32(rotated.map(value => Number(value)), [rotated.length, 1, 1, 1]);
+      using mask = flags.astype(Dtype.bool);
+      return ops.where(mask, restored, output);
+    };
+  }
+  makeEmptyBatch(): DelayedTurboQuantKVCache {
+    return new DelayedTurboQuantKVCache(this.kBits, this.vBits, this.start, this.fusedDecode, this.turboQuantKernels, undefined, this.compose);
+  }
+}

package/src/state/durability.ts

@@ -0,0 +1,214 @@
+import { cloneAttachments } from "./checkpoint";
+import { cleanupFailure, disposeResources } from "../runtime/resources";
+import type { Cache } from "../contracts/mlx/cache";
+import type { PromptCacheEntry } from "./prefix-cache";
+import type { SpillItem } from "./persistence-types";
+import type { SpillQueue } from "./persistence";
+
+export interface DurabilityPromptCache {
+  findExact(tokens: number[], ns?: string): PromptCacheEntry | null;
+}
+
+export interface DurabilitySnapshotStats {
+  pendingSnapshots: number;
+  pendingSpills: number;
+  pendingSpillBytes: number;
+  droppedSpills: number;
+  failedSpills: number;
+}
+
+export interface DurabilityFlushResult extends DurabilitySnapshotStats {
+  durable: boolean;
+  flushedSnapshots: number;
+  missingSnapshots: number;
+  elapsedMs: number;
+}
+
+interface DirtySnapshot {
+  key: string;
+  tokens: number[];
+  ns: string;
+}
+
+type StoreOutcome = "stored" | "missing" | "failed";
+
+/**
+ * Turns debounced prompt-cache puts into a durability boundary.
+ *
+ * A dirty record stays present until SpillQueue reports that its atomic store
+ * completed. Queue drops and write failures therefore remain retryable during
+ * an explicit flush or graceful shutdown.
+ */
+export class SsdDurabilityCoordinator {
+  readonly #dirty = new Map<string, DirtySnapshot>();
+  readonly #timers = new Map<string, ReturnType<typeof setTimeout>>();
+  readonly #attempts = new Map<string, Promise<StoreOutcome>>();
+  #flush: Promise<DurabilityFlushResult> | null = null;
+
+  constructor(
+    readonly promptCache: DurabilityPromptCache,
+    readonly spillQueue: SpillQueue,
+    readonly cloneCaches: (caches: Cache[]) => Cache[],
+    readonly isAlreadyDurable: (tokens: number[], ns: string) => boolean = () => false,
+    readonly debounceMs = 0,
+    readonly retryMs = 5_000,
+    readonly onStored: () => void = () => {},
+  ) {}
+
+  schedule(tokens: number[], ns = "", replace = true): void {
+    // Exact tensor objects use an empty history and still need persistence.
+    if (tokens.length === 0 && !this.promptCache.findExact(tokens, ns)?.attachments?.length) return;
+    // PromptCache can hold unrelated entries with the same namespace and
+    // length. Include the tokens so one conversation cannot cancel another
+    // conversation's pending durability record.
+    const key = `${ns.length}:${ns}:${tokens.join(",")}`;
+    if (!replace && this.#dirty.has(key)) return;
+    const rec: DirtySnapshot = { key, tokens: [...tokens], ns };
+    this.#dirty.set(key, rec);
+    this.#arm(key, this.debounceMs);
+  }
+
+  get stats(): DurabilitySnapshotStats {
+    return {
+      pendingSnapshots: this.#dirty.size,
+      pendingSpills: this.spillQueue.pendingCount,
+      pendingSpillBytes: this.spillQueue.pendingBytes,
+      droppedSpills: this.spillQueue.droppedCount,
+      failedSpills: this.spillQueue.failedCount,
+    };
+  }
+
+  /** Force every dirty RAM snapshot through the queue and wait for storage. */
+  flush(): Promise<DurabilityFlushResult> {
+    if (this.#flush) return this.#flush;
+    const started = performance.now();
+    this.#flush = this.#flushInner(started).finally(() => { this.#flush = null; });
+    return this.#flush;
+  }
+
+  #arm(key: string, delayMs: number): void {
+    const old = this.#timers.get(key);
+    if (old) clearTimeout(old);
+    const timer = setTimeout(() => {
+      this.#timers.delete(key);
+      void this.#attempt(key, false);
+    }, delayMs);
+    timer.unref?.();
+    this.#timers.set(key, timer);
+  }
+
+  async #attempt(key: string, force: boolean): Promise<StoreOutcome> {
+    const running = this.#attempts.get(key);
+    if (running) {
+      const outcome = await running;
+      if (!force || !this.#dirty.has(key)) return outcome;
+    }
+
+    const rec = this.#dirty.get(key);
+    if (!rec) return "stored";
+
+    const task = this.#store(rec);
+    this.#attempts.set(key, task);
+    try {
+      const outcome = await task;
+      if (outcome === "stored" && this.#dirty.get(key) === rec)
+        this.#dirty.delete(key);
+      if (outcome !== "stored" && !force && this.#dirty.get(key) === rec)
+        this.#arm(key, this.retryMs);
+      return outcome;
+    } finally {
+      if (this.#attempts.get(key) === task) this.#attempts.delete(key);
+    }
+  }
+
+  async #store(rec: DirtySnapshot): Promise<StoreOutcome> {
+    let snap: SpillItem | null = null;
+    try {
+      snap = (() => {
+        if (this.isAlreadyDurable(rec.tokens, rec.ns)) return null;
+        const entry = this.promptCache.findExact(rec.tokens, rec.ns);
+        if (!entry) return null;
+        const caches = this.cloneCaches(entry.caches);
+        try {
+          return { tokens: [...entry.tokens], caches, ns: rec.ns,
+            attachments: cloneAttachments(entry.attachments) };
+        } catch (error) {
+          return cleanupFailure(error, () => disposeResources(caches));
+        }
+      })();
+    } catch {
+      return "failed";
+    }
+    if (!snap)
+      return this.isAlreadyDurable(rec.tokens, rec.ns) ? "stored" : "missing";
+    const stored = await this.spillQueue.enqueue(snap);
+    // The queue releases its views before this continuation runs. Residency
+    // policy now sees memory without the completed writer's ownership.
+    if (stored) this.onStored();
+    return stored ? "stored" : "failed";
+  }
+
+  async #flushInner(started: number): Promise<DurabilityFlushResult> {
+    for (const timer of this.#timers.values()) clearTimeout(timer);
+    this.#timers.clear();
+
+    let flushedSnapshots = 0;
+    const missing = new Map<string, DirtySnapshot>();
+    const droppedBefore = this.spillQueue.droppedCount;
+    const failedBefore = this.spillQueue.failedCount;
+
+    if (this.#attempts.size > 0) {
+      const results = await Promise.allSettled([...this.#attempts.values()]);
+      flushedSnapshots += results.filter(result => result.status === "fulfilled" && result.value === "stored").length;
+    }
+    await this.spillQueue.drain();
+    // A background failure may have rearmed its retry while flush waited.
+    // Forced attempts below own retries now; no timer may outlive this flush.
+    for (const timer of this.#timers.values()) clearTimeout(timer);
+    this.#timers.clear();
+
+    // Flush dirty entries serially without a second queue of snapshots.
+    // Attempt each record version once. Missing state stays dirty, so a
+    // second flush cannot report success without a real durable snapshot.
+    const attempted = new Set<DirtySnapshot>();
+    while (this.#dirty.size > 0) {
+      const records = [...this.#dirty.entries()].filter(([, rec]) => !attempted.has(rec));
+      if (records.length === 0) break;
+      for (const [key, rec] of records) {
+        if (this.#dirty.get(key) !== rec) continue;
+        attempted.add(rec);
+        const outcome = await this.#attempt(key, true);
+        await this.spillQueue.drain();
+        if (outcome === "stored") {
+          flushedSnapshots++;
+        } else if (outcome === "missing") {
+          missing.set(key, rec);
+        }
+      }
+    }
+
+    // A later trimmable snapshot may now cover a superseded RAM ancestor.
+    // Reconcile against committed SSD coverage after all writes settle;
+    // never clear a missing prefix merely because another write succeeded.
+    for (const [key, rec] of missing) {
+      if (!this.isAlreadyDurable(rec.tokens, rec.ns)) continue;
+      if (this.#dirty.get(key) === rec) this.#dirty.delete(key);
+      missing.delete(key);
+      flushedSnapshots++;
+    }
+
+    const stats = this.stats;
+    return {
+      ...stats,
+      durable:
+        stats.pendingSnapshots === 0 &&
+        stats.pendingSpills === 0 &&
+        stats.droppedSpills === droppedBefore &&
+        stats.failedSpills === failedBefore &&
+        missing.size === 0,
+      flushedSnapshots,
+      missingSnapshots: missing.size,
+      elapsedMs: performance.now() - started,
+    };
+  }
+}

package/src/state/encoder-cache.ts

@@ -0,0 +1,26 @@
+import type { ObjectCache } from "../contracts/portable/object-cache";
+import { cloneAttachments, disposeAttachments } from "./checkpoint";
+import { type CheckpointAttachment } from "../contracts/mlx/checkpoint";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+
+/** Encoder-owned tensor schema over the cache's exact-object interface. */
+export class EncoderCache {
+  constructor(private readonly objects: ObjectCache<CheckpointAttachment[]>,
+    readonly identity: string) {}
+
+  async take(key: string): Promise<MlxArray | null> {
+    const lease = await this.objects.take(`encoder-v1:${this.identity}:${key}`);
+    if (!lease) return null;
+    try {
+      const tensor = lease.value[0]!.tensors[0]!;
+      return tensor.slice(tensor.shape.map(() => 0), [...tensor.shape]);
+    } finally { lease.dispose(); }
+  }
+
+  /** Borrow the result; storage retains its own immutable native view. */
+  put(key: string, features: MlxArray): void {
+    const owned = cloneAttachments([{ schema: "encoder-features-v1", metadata: {}, tensors: [features] }]);
+    try { this.objects.put(`encoder-v1:${this.identity}:${key}`, owned); }
+    catch (error) { disposeAttachments(owned); throw error; }
+  }
+}

package/src/state/full-kv-row-append.ts

@@ -0,0 +1,33 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache } from "../contracts/mlx/cache";
+import { fullRowPadding } from "./full-prefill-row";
+import { kvStorage } from "./kv-storage";
+import { disposeResources } from "../runtime/resources";
+
+/** Assemble full-attention tensors at the model's B while storage retains
+ * per-row formats and positions. Returned arrays own their captured state. */
+export function appendFullKvRows(rows: readonly Cache[], leftPad: readonly number[], width: number,
+  k: MlxArray, v: MlxArray, deferredValues = false): [MlxArray, MlxArray] {
+  const parts: [MlxArray, MlxArray][] = [], padded: MlxArray[] = [];
+  try {
+    for (const [index, row] of rows.map(kvStorage).entries()) {
+      using kr = k.slice([index, 0, 0, 0], [index + 1, k.shape[1]!, k.shape[2]!, k.shape[3]!]);
+      using vr = v.slice([index, 0, 0, 0], [index + 1, v.shape[1]!, v.shape[2]!, v.shape[3]!]);
+      parts.push(deferredValues && row.rotatedValueAttention
+        ? row.rotatedValueAttention.updateAndFetchDeferredV(kr, vr) : row.updateAndFetch(kr, vr));
+    }
+    const combine = (field: 0 | 1) => {
+      const views = parts.map((part, row) => {
+        const a = part[field], before = leftPad[row]! - fullRowPadding(rows[row]!), after = width - before - a.shape[2]!;
+        if (!before && !after) return a;
+        using prefix = ops.zeros([1, a.shape[1]!, before, a.shape[3]!], a.dtype);
+        using suffix = ops.zeros([1, a.shape[1]!, after, a.shape[3]!], a.dtype);
+        const result = ops.concatAxis([prefix, a, suffix], 2); padded.push(result); return result;
+      });
+      return views.length === 1 ? ops.contiguous(views[0]!) : ops.concatAxis(views, 0);
+    };
+    const keys = combine(0);
+    try { return [keys, combine(1)]; } catch (error) { keys.dispose(); throw error; }
+  } finally { disposeResources([...parts.flat(), ...padded]); }
+}

package/src/state/full-kv-row-donor.ts

@@ -0,0 +1,54 @@
+import { captureKvDonorAttention, combineKvDonorAttention } from "./kv-attention-view";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache, KvDonorRows } from "../contracts/mlx/cache";
+import { fullRowInner, fullRowPadding } from "./full-prefill-row";
+import { disposeResources } from "../runtime/resources";
+
+/** Read full-attention rows in their existing physical columns. Valid token
+ * coverage excludes prefill padding; no append, conversion or compaction. */
+export function captureFullKvDonorRows(rows: readonly Cache[], leftPad: readonly number[], width: number): KvDonorRows {
+  const parts: KvDonorRows[] = [], padded: MlxArray[] = [];
+  try {
+    for (const row of rows) parts.push(fullRowInner(row).captureDonorRows!());
+    const combine = (field: "keys" | "values") => {
+      const views = parts.map((part, index) => {
+        const a = part[field], before = leftPad[index]! - fullRowPadding(rows[index]!), after = width - before - a.shape[2]!;
+        if (!before && !after) return a;
+        using prefix = ops.zeros([1, a.shape[1]!, before, a.shape[3]!], a.dtype);
+        using suffix = ops.zeros([1, a.shape[1]!, after, a.shape[3]!], a.dtype);
+        const result = ops.concatAxis([prefix, a, suffix], 2); padded.push(result); return result;
+      });
+      return views.length === 1 ? ops.contiguous(views[0]!) : ops.concatAxis(views, 0);
+    };
+    const offsets = rows.map(row => row.offset), starts = [...leftPad];
+    const keys = combine("keys");
+    try { return { keys, values: combine("values"), offsets, starts, ends: offsets.map((offset, row) => starts[row]! + offset) }; }
+    catch (error) { keys.dispose(); throw error; }
+  } finally { disposeResources([...parts.flatMap(part => [part.keys, part.values]), ...padded]); }
+}
+
+
+/** Capture full rows in global columns while each encoded view retains its
+ * own physical width and numerical attention operation. */
+export function captureFullKvDonorAttention(rows: readonly Cache[], leftPad: readonly number[], width: number): import("../contracts/mlx/cache").KvDonorAttention {
+  const captured: import("../contracts/mlx/cache").KvDonorAttention[] = [];
+  try {
+    for (const [index, row] of rows.entries()) {
+      const view = captureKvDonorAttention(fullRowInner(row));
+      const pad = leftPad[index]! - fullRowPadding(row);
+      captured.push({ width, dtype: view.dtype, offsets: [row.offset],
+        starts: [leftPad[index]!], ends: [leftPad[index]! + row.offset],
+        attend(q, scale, mask) {
+          if (!mask.arr) return view.attend(q, scale, mask);
+          const from = mask.arr.shape.map(() => 0), to = [...mask.arr.shape];
+          from[from.length - 1] = pad; to[to.length - 1] = pad + view.width;
+          using arr = mask.arr.slice(from, to);
+          return view.attend(q, scale, { mode: mask.mode, arr });
+        },
+        dispose: () => view.dispose(),
+      });
+    }
+    return combineKvDonorAttention(captured);
+  } catch (error) { disposeResources(captured); throw error; }
+}

package/src/state/full-prefill-padding.ts

@@ -0,0 +1,56 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { PrefillPadding } from "../contracts/mlx/cache";
+import { rowRollIndices } from "./batched-row-storage";
+import { disposeResources } from "../runtime/resources";
+
+interface Positions { rowOffsets: number[]; leftPad: number[] }
+
+/** mlx-lm BatchKVCache.prepare/finalize (MIT), over arbitrary stored planes.
+ * Padding changes token positions; the representation's codec is uninvolved. */
+export class FullPrefillPadding {
+  constructor(readonly tokenAxis = 2) {}
+  #ends?: number[];
+  get endOffsets(): readonly number[] | undefined { return this.#ends; }
+  restoreEnds(ends: readonly number[] | undefined): void { this.#ends = ends ? [...ends] : undefined; }
+  validOffset(offset: number, row: number): number { return Math.min(offset, this.#ends?.[row] ?? offset); }
+
+  prepare(positions: Positions, padding: PrefillPadding): void {
+    if (!positions.rowOffsets.length) {
+      positions.rowOffsets = padding.lengths.map(() => 0);
+      positions.leftPad = [...positions.rowOffsets];
+    }
+    if (padding.leftPadding) {
+      positions.leftPad = positions.leftPad.map((pad, row) => pad + padding.leftPadding![row]!);
+      positions.rowOffsets = positions.rowOffsets.map((offset, row) => offset - padding.leftPadding![row]!);
+    }
+    // Like rotating storage, retain absolute valid ends. A cohort can finish
+    // before its planned width when a longer request leaves during prefill.
+    this.#ends = padding.rightPadding?.some(pad => pad > 0)
+      ? padding.lengths.map((length, row) => positions.rowOffsets[row]! + length) : undefined;
+  }
+
+  /** Borrow planes; return owned replacements only when padding needs moving.
+   * Stage every plane before committing positions, including on FFI failure. */
+  finalize(planes: readonly MlxArray[], positions: Positions): MlxArray[] | undefined {
+    if (!this.#ends) return;
+    const padding = this.#ends.map((end, row) => Math.max(0, positions.rowOffsets[row]! - end));
+    if (!padding.some(pad => pad > 0)) { this.#ends = undefined; return; }
+    const next: MlxArray[] = [];
+    try {
+      if (planes.length) {
+        const width = planes[0]!.shape[this.tokenAxis]!;
+        using indices = rowRollIndices(width, padding);
+        using compressed = this.tokenAxis === 1 ? ops.reshape(indices, [padding.length, width, 1]) : null;
+        for (const plane of planes) next.push(ops.takeAlongAxis(plane, compressed ?? indices, this.tokenAxis));
+      }
+    } catch (error) { disposeResources(next); throw error; }
+    positions.rowOffsets = positions.rowOffsets.map((offset, row) => offset - padding[row]!);
+    positions.leftPad = positions.leftPad.map((pad, row) => pad + padding[row]!);
+    this.#ends = undefined;
+    return next;
+  }
+
+  filter(keep: readonly number[]): void { if (this.#ends) this.#ends = keep.map(row => this.#ends![row]!); }
+  clear(): void { this.#ends = undefined; }
+}

package/src/state/full-prefill-row.ts

@@ -0,0 +1,112 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { materializeCopy } from "@mlx-bun/mlx/materialize";
+import { withResource, disposeResources } from "../runtime/resources";
+import { leaseCacheState } from "./leases";
+import { FullPrefillPadding } from "./full-prefill-padding";
+import { KVCache, type AffineKernelSet } from "./kv";
+import { QuantizedKVCache } from "./quantized-kv";
+import { TurboQuantKVCache } from "./turboquant-kv";
+import { type Cache, type Mask, type PrefillPadding } from "../contracts/mlx/cache";
+import { kvStorage } from "./kv-storage";
+
+type FullCache = KVCache | QuantizedKVCache | TurboQuantKVCache;
+
+/** A physical full-attention row with separate valid-token coverage. Codecs
+ * consume the same physical tensors; padding never advances their threshold. */
+export class FullPrefillRow implements Cache {
+  readonly padding = new FullPrefillPadding();
+  leftPadding = 0;
+  constructor(public inner: FullCache) {}
+  get minimumReusableOffset(): number { return (this.inner as Cache).minimumReusableOffset ?? 0; }
+  set minimumReusableOffset(value: number) { (this.inner as Cache).minimumReusableOffset = value; }
+  get positionOffset(): number { return this.inner.offset - this.leftPadding; }
+  get physicalLength(): number { return this.inner.offset; }
+  get offset(): number { return Math.max(0, this.padding.validOffset(this.positionOffset, 0)); }
+  get stateNeedsDispose(): boolean { return (this.inner as Cache).stateNeedsDispose ?? false; }
+  get affineConversion(): FullPrefillRow | undefined { return this.inner instanceof KVCache ? this : undefined; }
+  get turboConversion(): FullPrefillRow | undefined { return this.inner instanceof KVCache ? this : undefined; }
+  get quantizedAttention() { return kvStorage(this.inner).quantizedAttention; }
+  get rotatedValueAttention() { return kvStorage(this.inner).rotatedValueAttention; }
+  signature(): string { return this.inner.signature(); }
+  state(): MlxArray[] { return this.inner.state(); }
+  makeMask(tokens: number, window: number | null): Mask { return this.inner.makeMask(tokens, window); }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void { this.inner.trim(count); }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.inner.updateAndFetch(k, v); }
+  bytesPerToken(): number { return (this.inner as Cache).bytesPerToken?.() ?? 0; }
+  preparePrefill(padding: PrefillPadding): void {
+    const positions = { rowOffsets: [this.positionOffset], leftPad: [this.leftPadding] };
+    this.padding.prepare(positions, padding);
+    this.leftPadding = positions.leftPad[0]!;
+  }
+  finalizePrefill(): void {
+    const positions = { rowOffsets: [this.positionOffset], leftPad: [this.leftPadding] };
+    const next = withResource(leaseCacheState(this.inner), state => this.padding.finalize(state, positions));
+    if (next) this.replace(next, this.inner.offset);
+    this.leftPadding = positions.leftPad[0]!;
+  }
+  removeLeftPadding(count: number): void {
+    if (!count) return;
+    const next: MlxArray[] = [];
+    try {
+      withResource(leaseCacheState(this.inner), state => {
+        for (const plane of state) next.push(plane.slice([0, 0, count, 0], [...plane.shape]));
+      });
+      this.replace(next, this.inner.offset - count); next.length = 0;
+      this.leftPadding -= count;
+    } finally { disposeResources(next); }
+  }
+  private restored(planes: MlxArray[], offset: number): FullCache {
+    const source = this.inner;
+    if (source instanceof KVCache) {
+      const result = new KVCache(undefined, source.turboQuantKernels, source.windowRead);
+      if (planes.length) result.restoreState(planes[0]!, planes[1]!, offset);
+      return result;
+    }
+    if (source instanceof QuantizedKVCache) {
+      const result = new QuantizedKVCache(source.groupSize, source.bits, source.kernels);
+      if (planes.length) result.restoreState({ packed: planes[0]!, scales: planes[1]!, biases: planes[2]! },
+        { packed: planes[3]!, scales: planes[4]!, biases: planes[5]! }, offset);
+      return result;
+    }
+    const result = new TurboQuantKVCache(source.kBits, source.vBits, source.turboQuantKernels, source.fusedDecode);
+    if (planes.length) result.restoreState({ kIdx: planes[0]!, kScales: planes[1]!, kZeros: planes[2]!,
+      vPacked: planes[3]!, vScales: planes[4]! }, offset, source.headDim!);
+    return result;
+  }
+  private replace(planes: MlxArray[], offset: number): void {
+    const next = this.restored(planes, offset), previous = this.inner;
+    (next as Cache).minimumReusableOffset = this.minimumReusableOffset;
+    this.inner = next; previous.dispose();
+  }
+  extract(): Cache {
+    const count = this.offset, planes: MlxArray[] = [];
+    try {
+      if (count) withResource(leaseCacheState(this.inner), state => {
+        for (const plane of state) {
+          using view = plane.slice([0, 0, this.leftPadding, 0], [1, plane.shape[1]!, this.leftPadding + count, plane.shape[3]!]);
+          planes.push(materializeCopy(view));
+        }
+      });
+      const result = this.restored(planes, count); planes.length = 0;
+      (result as Cache).minimumReusableOffset = this.minimumReusableOffset;
+      return result;
+    } finally { disposeResources(planes); }
+  }
+  /** The plain inner storage's conversion (`KVCache.toQuantized`): the
+   * converted row attends with the set `kernels` builds. */
+  toQuantized(groupSize: number, bits: number, kernels: AffineKernelSet): FullPrefillRow {
+    this.inner = (this.inner as KVCache).toQuantized(groupSize, bits, kernels);
+    return this;
+  }
+  toTurboQuantized(kBits: number, vBits: number, fusedDecode?: boolean): FullPrefillRow {
+    const source = this.inner as KVCache;
+    this.inner = TurboQuantKVCache.fromKVCache(source, kBits, vBits, source.turboQuantKernels, fusedDecode);
+    return this;
+  }
+  dispose(): void { this.inner.dispose(); this.padding.clear(); }
+}
+
+export function fullRowInner(row: Cache): Cache { return row instanceof FullPrefillRow ? row.inner : row; }
+export function fullRowPadding(row: Cache): number { return row instanceof FullPrefillRow ? row.leftPadding : 0; }
+export function fullRowPhysicalLength(row: Cache): number { return row instanceof FullPrefillRow ? row.physicalLength : row.offset; }