mlx-bun 0.3.0 → 0.4.0

package/src/model/rotating-kv-layout.ts

@@ -0,0 +1,262 @@
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { cloneKvCaches } from "../kv-store";
+import { type BatchableCache, type Cache, type Mask, type PaddedPrefillCache, type PrefillPadding, RotatingKVCache, RotatingQuantizedKVCache } from "./gemma4-base";
+import { BatchedRotatingCache, buildBatchedRotatingMask } from "./batched-rotating";
+import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
+import { BatchedRotatingState, type RotatingPositionSnapshot } from "./batched-rotating-state";
+import { plainRowStorage, quantizedRowStorage, temporalStorageView } from "./batched-row-storage";
+import type { TransitioningKvPositions, TransitionedKvLayout } from "../backends/mlx/transitioning-kv-rows";
+import { captureRotatingDonorAttention } from "./rotating-kv-donor";
+
+import { rollbackRotatingRing } from "./rotating-row-transaction";
+
+const fields = ["packed", "scales", "biases"] as const;
+type Ring = BatchedRotatingCache | BatchedRotatingQuantCache;
+
+function combinedPosition(rings: readonly Ring[]): RotatingPositionSnapshot {
+  const states = rings.map(ring => ring.positionSnapshot);
+  return { ...states[0]!, offsets: states.map(state => state.offsets[0]!),
+    leftPad: states.map(state => state.leftPad[0]!),
+    prefillEnds: states[0]!.prefillEnds
+      ? states.map(state => state.prefillEnds![0]!) : undefined };
+}
+
+/** One logical request with the group's physical ring geometry. This keeps
+ * masks stable while different requests cross their precision boundary. */
+export class AlignedRotatingCache implements Cache {
+  minimumReusableOffset = 0;
+  constructor(public inner: Ring) {}
+  get offset(): number { return this.inner.validOffset(0); }
+  get affineConversion(): AlignedRotatingCache | undefined { return this.inner instanceof BatchedRotatingCache ? this : undefined; }
+  get quantizedAttention() { return this.inner instanceof BatchedRotatingQuantCache ? this.inner : undefined; }
+  signature(): string { return this.inner.signature(); }
+  captureDonorAttention() { return captureRotatingDonorAttention(this.inner); }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.inner.updateAndFetch(k, v); }
+  makeMask(tokens: number, window: number | null): Mask { return this.inner.makeMask(tokens, window); }
+  isTrimmable(): boolean { return this.inner.positionSnapshot.totalOffset < this.inner.maxSize; }
+  trim(count: number): void { this.inner.trim(count); }
+  rollback(before: number, keep: number, preserveColumns: boolean): void {
+    this.inner = rollbackRotatingRing(this.inner, [before], keep, preserveColumns);
+  }
+  state(): MlxArray[] { return this.inner.state(); }
+  bytesPerToken(): number {
+    const state = this.state(), tokens = state[0]?.shape[2] ?? 0;
+    return tokens ? state.reduce((n, a) => n + a.nbytes, 0) / tokens : 0;
+  }
+  dispose(): void { this.inner.dispose(); }
+  extract(limit?: number): Cache {
+    const row = this.inner.extractRow(0, limit) ?? (this.inner instanceof BatchedRotatingQuantCache
+      ? new RotatingQuantizedKVCache(this.inner.maxSize, this.inner.groupSize, this.inner.bits)
+      : new RotatingKVCache(this.inner.maxSize));
+    row.minimumReusableOffset = this.minimumReusableOffset;
+    return row;
+  }
+  toQuantized(groupSize: number, bits: number): AlignedRotatingCache {
+    const plain = this.inner as BatchedRotatingCache;
+    const keys = ops.quantize(plain.keys!, groupSize, bits);
+    let values: ops.QuantizedTensor;
+    try { values = ops.quantize(plain.values!, groupSize, bits); }
+    catch (error) { quantizedRowStorage.dispose(keys); throw error; }
+    const next = new AlignedRotatingCache(BatchedRotatingQuantCache.adoptPhysical(keys, values, groupSize, bits, plain.positionSnapshot));
+    next.minimumReusableOffset = this.minimumReusableOffset;
+    this.dispose();
+    return next;
+  }
+}
+
+/** Borrow the newest valid window from a solo cache, including a temporary
+ * oversized block. No dequantization or physical-ring normalization in place. */
+export function rotatingSourcePosition(row: RotatingKVCache | RotatingQuantizedKVCache): BatchedRotatingState {
+  const state = new BatchedRotatingState(row.maxSize, [0], [row.offset]);
+  const length = row instanceof RotatingKVCache ? row.keys?.shape[2] ?? 0 : row.keys?.packed.shape[2] ?? 0;
+  state.restore({ maxSize: row.maxSize, offsets: [row.offset], leftPad: [0], totalOffset: row.offset,
+    ringIndex: row.ringIdx, rotated: row.ringIdx < row.offset && row.ringIdx < length });
+  return state;
+}
+
+/** Take ownership of solo rows and align them once at membership changes. */
+export function alignRotatingRows(rows: Cache[]): AlignedRotatingCache[] {
+  const width = Math.max(0, ...rows.map(row => Math.min(row.offset, (row as RotatingKVCache).maxSize)));
+  const prototype = rows.find(row => row.state().length > 0) as RotatingKVCache | RotatingQuantizedKVCache | undefined;
+  const emptyPlane = (field: "keys" | "values"): MlxArray | null => {
+    if (!width || !prototype) return null;
+    const tensor = prototype[field]!;
+    if (tensor instanceof MlxArray) return ops.zeros([1, tensor.shape[1]!, width, tensor.shape[3]!], tensor.dtype);
+    return ops.zeros([1, tensor.packed.shape[1]!, width, tensor.packed.shape[3]! * 32 / (prototype as RotatingQuantizedKVCache).bits], tensor.scales.dtype);
+  };
+  const result: AlignedRotatingCache[] = [];
+  try {
+    for (const source of rows) {
+      if (!(source instanceof RotatingKVCache || source instanceof RotatingQuantizedKVCache))
+        throw new Error(`rotating layout cannot adopt ${source.signature()}`);
+      const length = Math.min(source.offset, source.maxSize), pad = width - length;
+      const position: RotatingPositionSnapshot = { maxSize: source.maxSize, offsets: [source.offset], leftPad: [pad],
+        totalOffset: width, ringIndex: width, rotated: false };
+      const state = rotatingSourcePosition(source);
+      const options = { from: Math.max(0, state.activeLength - length), to: state.activeLength };
+      let inner: Ring;
+      if (source instanceof RotatingKVCache) {
+        using k = source.keys ? temporalStorageView(plainRowStorage, source.keys, state, options) : null;
+        using v = source.values ? temporalStorageView(plainRowStorage, source.values, state, options) : null;
+        const keys = k ? plainRowStorage.padLeft(k, pad) : emptyPlane("keys");
+        let values: MlxArray | null;
+        try { values = v ? plainRowStorage.padLeft(v, pad) : emptyPlane("values"); }
+        catch (error) { keys?.dispose(); throw error; }
+        inner = BatchedRotatingCache.adoptPhysical(keys, values, position);
+      } else {
+        const k = source.keys ? temporalStorageView(quantizedRowStorage, source.keys, state, options) : null;
+        const v = source.values ? temporalStorageView(quantizedRowStorage, source.values, state, options) : null;
+        const empty = (field: "keys" | "values") => {
+          using tensor = emptyPlane(field);
+          return tensor ? ops.quantize(tensor, source.groupSize, source.bits) : null;
+        };
+        try {
+          const keys = k ? quantizedRowStorage.padLeft(k, pad) : empty("keys");
+          let values: ops.QuantizedTensor | null;
+          try { values = v ? quantizedRowStorage.padLeft(v, pad) : empty("values"); }
+          catch (error) { if (keys) quantizedRowStorage.dispose(keys); throw error; }
+          inner = BatchedRotatingQuantCache.adoptPhysical(keys, values, source.groupSize, source.bits, position);
+        } finally { if (k) quantizedRowStorage.dispose(k); if (v) quantizedRowStorage.dispose(v); }
+      }
+      const row = new AlignedRotatingCache(inner); row.minimumReusableOffset = source.minimumReusableOffset ?? 0;
+      result.push(row);
+    }
+  } catch (error) { for (const row of result) row.dispose(); throw error; }
+  for (const row of rows) row.dispose();
+  return result;
+}
+
+/** Position state has no codec or attention arithmetic. */
+export class RotatingKvPositions implements TransitioningKvPositions {
+  #state: BatchedRotatingState;
+  #rope?: MlxArray;
+  constructor(maxSize: number) { this.#state = new BatchedRotatingState(maxSize, []); }
+  get rowOffsets(): readonly number[] { return this.#state.offsets; }
+  get leftPad(): readonly number[] { return this.#state.leftPad; }
+  get offset(): number { return this.#state.totalOffset; }
+  get batchSize(): number | null { return this.#state.batchSize || null; }
+  get ropeOffsetArr(): MlxArray | undefined {
+    if (!this.batchSize) return undefined;
+    return this.#rope ??= ops.fromInt32([...this.rowOffsets], [this.batchSize]);
+  }
+  sync(rows: readonly Cache[]): void {
+    this.#rope?.dispose(); this.#rope = undefined;
+    if (!rows.length) { this.#state = new BatchedRotatingState(this.#state.maxSize, []); return; }
+    this.#state.restore(combinedPosition(rows.map(row => (row as AlignedRotatingCache).inner)));
+  }
+  filterRows(keep: readonly number[]): void { this.#state.filter(keep); this.#rope?.dispose(); this.#rope = undefined; }
+  makeMask(tokens: number, window: number | null): Mask {
+    return { mode: "array", arr: buildBatchedRotatingMask(this.#state.batchSize, tokens, this.#state.leftPad,
+      this.#state.maxSize, window ?? this.#state.maxSize, this.#state.ringIndex, this.#state.totalOffset, this.#state.rotated,
+      tokens > 1 || this.#state.hasPendingPadding) };
+  }
+  dispose(): void { this.#rope?.dispose(); this.#rope = undefined; this.#state = new BatchedRotatingState(this.#state.maxSize, []); }
+}
+
+/** Once all rows use affine storage, concatenate their already-aligned planes
+ * without changing ring columns or rebuilding any attention mask. */
+export class RotatingAffineLayout implements TransitionedKvLayout, PaddedPrefillCache {
+  #inner?: BatchedRotatingQuantCache;
+  #before?: number[];
+  #minimum: number[] = [];
+  constructor(readonly maxSize: number, readonly groupSize: number, readonly bits: number) {}
+  get minimumReusableOffset(): number { return Math.max(0, ...this.#minimum); }
+  get quantizedAttention() { return this.#inner; }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    return captureRotatingDonorAttention(this.#inner!);
+  }
+  get rowOffsets(): readonly number[] { return this.#inner?.offsetArr ?? []; }
+  get leftPad(): readonly number[] { return this.#inner?.leftPad ?? []; }
+  get offset(): number { return this.#inner?.offset ?? 0; }
+  get batchSize(): number | null { return this.#inner?.batchSize || null; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.#inner?.ropeOffsetArr; }
+  get maxTokens(): number { return this.maxSize; }
+  signature(): string { return `kv:rotating-affine:${this.maxSize}:${this.groupSize}:${this.bits}`; }
+  makeEmptyBatch(): RotatingAffineLayout { return new RotatingAffineLayout(this.maxSize, this.groupSize, this.bits); }
+  makeMask(tokens: number, window: number | null): Mask { return this.#inner!.makeMask(tokens, window); }
+  preparePrefill(padding: PrefillPadding): void { this.#inner!.preparePrefill(padding); }
+  finalizePrefill(): void { this.#inner!.finalizePrefill(); }
+  bytesPerToken(): number {
+    const tokens = this.#inner?.keys?.packed.shape[2] ?? 0;
+    return tokens ? this.state().reduce((n, a) => n + a.nbytes, 0) / (tokens * this.#inner!.batchSize) : 0;
+  }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
+  isTrimmable(): boolean { return this.#inner!.positionSnapshot.totalOffset < this.maxSize; }
+  trim(count: number): void { this.#inner!.trim(count); }
+  specRoundBegin(): void { this.#before = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#before = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    this.#inner = rollbackRotatingRing(this.#inner!, this.#before!, keep);
+    this.#before = undefined;
+  }
+  updateAndFetch(): [MlxArray, MlxArray] { return this.#inner!.updateAndFetch(); }
+  releaseRopeArr(): void { this.#inner?.releaseRopeArr(); }
+  mergeRows(rows: readonly Cache[]): void {
+    const owned: Cache[] = [];
+    try {
+      for (const row of rows) {
+        if (row instanceof RotatingAffineLayout) {
+          for (let i = 0; i < (row.batchSize ?? 0); i++) owned.push(row.extractRow(i));
+        } else if (row instanceof AlignedRotatingCache) owned.push(row.extract());
+        else owned.push(...cloneKvCaches([row]));
+      }
+      if (!owned.length) { this.dispose(); return; }
+      const aligned = alignRotatingRows(owned); owned.length = 0;
+      try { this.adoptAlignedRows(aligned); } finally { for (const row of aligned) row.dispose(); }
+    } finally { for (const row of owned) row.dispose(); }
+  }
+  adoptAlignedRows(rows: readonly AlignedRotatingCache[]): void {
+    const rings = rows.map(row => (row as AlignedRotatingCache).inner as BatchedRotatingQuantCache);
+    const position = combinedPosition(rings);
+    if (rings.every(row => !row.keys)) {
+      const empty = BatchedRotatingQuantCache.adoptPhysical(null, null, this.groupSize, this.bits,
+        position);
+      this.#inner?.dispose(); this.#inner = empty; this.#minimum = rows.map(row => row.minimumReusableOffset ?? 0);
+      return;
+    }
+    const merge = (field: "keys" | "values") => {
+      const planes: MlxArray[] = [];
+      try {
+        for (const part of fields) planes.push(ops.concatAxis(rings.map(row => row[field]![part]), 0));
+        return { packed: planes[0]!, scales: planes[1]!, biases: planes[2]! };
+      } catch (error) { for (const plane of planes) plane.dispose(); throw error; }
+    };
+    const keys = merge("keys"); let values: ops.QuantizedTensor;
+    try { values = merge("values"); } catch (error) { quantizedRowStorage.dispose(keys); throw error; }
+    const inner = BatchedRotatingQuantCache.adoptPhysical(keys, values, this.groupSize, this.bits,
+      position);
+    this.#inner?.dispose(); this.#inner = inner; this.#minimum = rows.map(row => row.minimumReusableOffset ?? 0);
+  }
+  filterRows(keep: readonly number[]): void { this.#inner!.filterRows(keep); this.#minimum = keep.map(row => this.#minimum[row]!); }
+  extractRow(row: number): Cache {
+    return this.extract(row);
+  }
+  protected extract(row: number, limit?: number): Cache {
+    const result = this.#inner!.extractRow(row, limit) ?? new RotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits);
+    result.minimumReusableOffset = this.#minimum[row] ?? 0; return result;
+  }
+  state(): MlxArray[] { return this.#inner?.state() ?? []; }
+  dispose(): void { this.#inner?.dispose(); this.#inner = undefined; this.#minimum = []; this.#before = undefined; }
+}
+
+/** A committed speculative checkpoint contains only the newest valid window. */
+export class SpeculativeRotatingAffineLayout extends RotatingAffineLayout {
+  #ropeForOffset = -1;
+  override get ropeOffsetArr(): MlxArray | undefined {
+    if (this.#ropeForOffset !== this.offset) {
+      this.releaseRopeArr(); this.#ropeForOffset = this.offset;
+    }
+    return super.ropeOffsetArr;
+  }
+  override makeEmptyBatch(): SpeculativeRotatingAffineLayout {
+    return new SpeculativeRotatingAffineLayout(this.maxSize, this.groupSize, this.bits);
+  }
+  override makeMask(tokens: number, window: number | null): Mask {
+    if (this.leftPad.every(pad => pad === 0) &&
+        this.offset + tokens <= Math.min(this.maxSize, window ?? this.maxSize))
+      return { mode: tokens === 1 ? "" : "causal", arr: null };
+    return super.makeMask(tokens, window);
+  }
+  override extractRow(row: number): Cache { return this.extract(row, this.maxSize); }
+}

package/src/model/rotating-row-transaction.ts

@@ -0,0 +1,35 @@
+import * as ops from "../mlx/ops";
+import { BatchedRotatingCache } from "./batched-rotating";
+import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
+import { rotatingRollbackPlan } from "./batched-rotating-state";
+import { plainRowStorage, quantizedRowStorage, rowRollIndices } from "./batched-row-storage";
+
+export type RotatingRing = BatchedRotatingCache | BatchedRotatingQuantCache;
+
+/** Retain accepted columns without decoding or recomputing the ring. Ownership
+ * moves to the returned ring; aligned members preserve their shared rectangle. */
+export function rollbackRotatingRing(ring: BatchedRotatingCache, before: readonly number[], keep: number | readonly number[], preserveColumns?: boolean): BatchedRotatingCache;
+export function rollbackRotatingRing(ring: BatchedRotatingQuantCache, before: readonly number[], keep: number | readonly number[], preserveColumns?: boolean): BatchedRotatingQuantCache;
+export function rollbackRotatingRing(ring: RotatingRing, before: readonly number[], keep: number | readonly number[], preserveColumns?: boolean): RotatingRing;
+export function rollbackRotatingRing(ring: RotatingRing, before: readonly number[],
+  keep: number | readonly number[], preserveColumns = false): RotatingRing {
+  const plan = rotatingRollbackPlan(ring.positionSnapshot, before, keep, preserveColumns);
+  if ("trim" in plan) { ring.trim(plan.trim); ring.releaseRopeArr(); return ring; }
+  let next: RotatingRing;
+  if (ring instanceof BatchedRotatingQuantCache) {
+    using indices = rowRollIndices(ring.keys!.packed.shape[2]!, plan.shifts);
+    const keys = quantizedRowStorage.rollRows(ring.keys!, indices);
+    let values: ops.QuantizedTensor;
+    try { values = quantizedRowStorage.rollRows(ring.values!, indices); }
+    catch (error) { quantizedRowStorage.dispose(keys); throw error; }
+    next = BatchedRotatingQuantCache.adoptPhysical(keys, values, ring.groupSize, ring.bits, plan.position);
+  } else {
+    using indices = rowRollIndices(ring.keys!.shape[2]!, plan.shifts);
+    const keys = plainRowStorage.rollRows(ring.keys!, indices);
+    let values: import("../mlx/array").MlxArray;
+    try { values = plainRowStorage.rollRows(ring.values!, indices); }
+    catch (error) { keys.dispose(); throw error; }
+    next = BatchedRotatingCache.adoptPhysical(keys, values, plan.position);
+  }
+  ring.dispose(); return next;
+}

package/src/model/speculative-rotating-kv.ts

@@ -0,0 +1,101 @@
+import { rotatingSourcePosition } from "./rotating-kv-layout";
+import { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { type BatchableCache, type Cache, type Mask, type PrefillPadding, RotatingKVCache } from "./gemma4-base";
+import { BatchedRotatingCache } from "./batched-rotating";
+import { rollbackRotatingRing } from "./rotating-row-transaction";
+import { BatchedRotatingState } from "./batched-rotating-state";
+import { plainRowStorage, temporalStorageView } from "./batched-row-storage";
+import { disposeResources } from "../engine/resources";
+
+/** Ring transactions retain accepted KV columns, including rounds that cross
+ * the window boundary. The verify block already retains the complete history
+ * needed by its earliest query; no target recomputation or scheduler policy. */
+export class SpeculativeRotatingKVCache implements BatchableCache {
+  #inner: BatchedRotatingCache;
+  #before?: number[];
+  constructor(readonly maxSize: number) { this.#inner = new BatchedRotatingCache(maxSize, []); }
+  signature(): string { return "kv:rotating-plain"; }
+  get offset(): number { return this.#inner.offset; }
+  get rowOffsets(): readonly number[] { return this.#inner.offsetArr; }
+  get leftPad(): readonly number[] { return this.#inner.leftPad; }
+  get batchSize(): number { return this.#inner.batchSize; }
+  get ropeOffsetArr(): MlxArray { return this.#inner.ropeOffsetArr; }
+  makeEmptyBatch(): SpeculativeRotatingKVCache { return new SpeculativeRotatingKVCache(this.maxSize); }
+  makeMask(tokens: number, window: number | null): Mask { return this.#inner.makeMask(tokens, window); }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#inner.updateAndFetch(k, v); }
+  preparePrefill(padding: PrefillPadding): void { this.#inner.preparePrefill(padding); }
+  finalizePrefill(): void { this.#inner.finalizePrefill(); }
+  state(): MlxArray[] { return this.#inner.state(); }
+  isTrimmable(): boolean { return this.#inner.isTrimmable(); }
+  trim(count: number): void { this.#inner.trim(count); }
+  bytesPerToken(): number {
+    const state = this.state();
+    return state.length ? state.reduce((n, a) => n + a.nbytes, 0) / (this.batchSize * state[0]!.shape[2]!) : 0;
+  }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
+  extractRow(row: number): Cache {
+    const result = new RotatingKVCache(this.maxSize);
+    if (!this.rowOffsets[row] || !this.#inner.keys || !this.#inner.values) return result;
+    const position = new BatchedRotatingState(this.maxSize, []); position.restore(this.#inner.positionSnapshot);
+    const range = { row, from: Math.max(0, position.leftPad[row]!, position.activeLength - this.maxSize),
+      to: position.activeLength, copy: true };
+    const keys = temporalStorageView(plainRowStorage, this.#inner.keys, position, range);
+    let values: MlxArray;
+    try { values = temporalStorageView(plainRowStorage, this.#inner.values, position, range); }
+    catch (error) { keys.dispose(); throw error; }
+    result.restoreState(keys, values, this.rowOffsets[row]!, keys.shape[2]!); return result;
+  }
+  captureDonorRows(): import("./gemma4-base").KvDonorRows {
+    const position = new BatchedRotatingState(this.maxSize, []); position.restore(this.#inner.positionSnapshot);
+    const from = Math.max(0, position.activeLength - this.maxSize);
+    const range = { from, to: position.activeLength };
+    const keys = temporalStorageView(plainRowStorage, this.#inner.keys!, position, range); let values: MlxArray;
+    try { values = temporalStorageView(plainRowStorage, this.#inner.values!, position, range); }
+    catch (error) { keys.dispose(); throw error; }
+    return { keys, values, offsets: [...this.rowOffsets], starts: this.leftPad.map(pad => Math.max(0, pad - from)),
+      ends: this.rowOffsets.map(() => position.activeLength - from) };
+  }
+
+  filterRows(keep: readonly number[]): void { this.#inner.filterRows(keep); }
+  mergeRows(sources: readonly Cache[]): void {
+    const only = sources.length === 1 ? sources[0] : undefined;
+    if (only instanceof SpeculativeRotatingKVCache) {
+      const keys = only.#inner.keys ? ops.copyOf(only.#inner.keys) : null;
+      let values: MlxArray | null;
+      try { values = only.#inner.values ? ops.copyOf(only.#inner.values) : null; }
+      catch (error) { keys?.dispose(); throw error; }
+      const next = BatchedRotatingCache.adoptPhysical(keys, values, only.#inner.positionSnapshot);
+      this.#inner.dispose(); this.#inner = next; return;
+    }
+    const held: Cache[] = [], planes: { keys: MlxArray; values: MlxArray }[] = [], offsets: number[] = [];
+    try {
+      const rows = sources.flatMap(source => source instanceof SpeculativeRotatingKVCache
+        ? Array.from({ length: source.batchSize }, (_, row) => { const cache = source.extractRow(row); held.push(cache); return cache; })
+        : [source]) as RotatingKVCache[];
+      const prototype = rows.find(row => row.keys !== null);
+      if (!prototype) {
+        this.#inner.dispose(); this.#inner = new BatchedRotatingCache(this.maxSize, rows.map(() => 0)); return;
+      }
+      for (const cache of rows) {
+        const state = rotatingSourcePosition(cache);
+        const range = { from: Math.max(0, state.activeLength - this.maxSize), to: state.activeLength };
+        const plane = (field: "keys" | "values") => cache[field]
+          ? temporalStorageView(plainRowStorage, cache[field]!, state, range)
+          : ops.zeros([1, prototype[field]!.shape[1]!, 0, prototype[field]!.shape[3]!], prototype[field]!.dtype);
+        const keys = plane("keys"); let values: MlxArray;
+        try { values = plane("values"); } catch (error) { keys.dispose(); throw error; }
+        planes.push({ keys, values }); offsets.push(cache.offset);
+      }
+      const next = BatchedRotatingCache.merge(planes, offsets, this.maxSize);
+      this.#inner.dispose(); this.#inner = next;
+    } finally { disposeResources([...held, ...planes.flatMap(row => [row.keys, row.values])]); }
+  }
+  specRoundBegin(): void { this.#before = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#before = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    this.#inner = rollbackRotatingRing(this.#inner, this.#before!, keep);
+    this.#before = undefined;
+  }
+  dispose(): void { this.#before = undefined; this.#inner.dispose(); }
+}

package/src/model/ssm-prefill-padding.ts

@@ -0,0 +1,39 @@
+import type { PrefillPadding } from "./gemma4-base";
+import type { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+
+/** ArraysCache's remaining lengths and padding from mlx-lm (MIT).
+ * Recurrent state has no token axis to roll: masked tokens must not update it. */
+export class SsmPrefillPadding {
+  #left: number[];
+  #lengths: number[];
+  constructor(padding: PrefillPadding) {
+    this.#left = padding.leftPadding ? [...padding.leftPadding] : padding.lengths.map(() => 0);
+    this.#lengths = [...padding.lengths];
+  }
+  convolutionLengths(tokens: number): number[] {
+    return this.#lengths.map(length => Math.max(0, Math.min(tokens, length)));
+  }
+  makeMask(tokens: number): MlxArray {
+    const B = this.#lengths.length;
+    using arange = ops.arange(0, tokens, 1, Dtype.int32);
+    using positions = ops.reshape(arange, [1, tokens]);
+    using left = ops.fromInt32(this.#left, [B, 1]);
+    using ends = ops.fromInt32(this.#lengths, [B, 1]);
+    using after = ops.greaterEqual(positions, left);
+    using before = ops.less(positions, ends);
+    return ops.logicalAnd(after, before);
+  }
+  advance(tokens: number): number[] {
+    const counts = this.#lengths.map((length, row) => Math.max(0,
+      Math.min(tokens, length) - Math.max(0, Math.min(tokens, this.#left[row]!))));
+    this.#lengths = this.#lengths.map(length => length - tokens);
+    this.#left = this.#left.map(left => left - tokens);
+    return counts;
+  }
+  filter(keep: readonly number[]): void {
+    this.#left = keep.map(row => this.#left[row]!);
+    this.#lengths = keep.map(row => this.#lengths[row]!);
+  }
+}

package/src/model/trellis-linear.ts

@@ -403,7 +403,7 @@ function lutFor(L: number): MlxArray {
   return a;
 }
 
-/** Decode variant (see HEADER). Default is the fastest measured on M1 Max;
+/** Decode variant (see HEADER). Default passed the M4 Pro closeout matrix;
  *  `MLX_BUN_TRELLIS_VARIANT` overrides for experiments. */
 let variantOverride: number | null = null;
 /** Explicit override (benches); otherwise the runtime flag — read per call so
@@ -411,7 +411,7 @@ let variantOverride: number | null = null;
  *  can A/B two decodes on one loaded model. */
 export function setTrellisVariant(v: number | null): void { variantOverride = v; }
 function variant(): number {
-  return variantOverride ?? runtimeNumber("MLX_BUN_TRELLIS_VARIANT", 6);
+  return variantOverride ?? runtimeNumber("MLX_BUN_TRELLIS_VARIANT", 13);
 }
 // Variants 7..13 change work assignment/scheduling, not decoded weight values.
 function decoderVariant(selected = variant()): number { return selected >= 7 && selected <= 13 ? 6 : selected; }

package/src/model/turboquant-codec.ts

@@ -0,0 +1,83 @@
+import type { MlxArray } from "../mlx/array";
+import * as tq from "../mlx/turboquant-ops";
+import type { KvCodec } from "../backends/mlx/kv-codec";
+
+/** One TurboQuant-encoded (K, V) storage tuple — the 5 arrays kv-store.ts
+ *  and toQuantized/fromKVCache pass around. Not `ops.QuantizedTensor`
+ *  (mlx's affine int4/int8 scheme): this is the rotation + Lloyd-Max
+ *  layout (docs/design/turboquant.md), asymmetric-affine for keys,
+ *  FWHT+Lloyd-Max for values. Field order is the kv-store.ts tensor-slot
+ *  contract — do not reorder without updating snapshotCache/loadKvCache. */
+export interface TurboQuantTensor {
+  kIdx: MlxArray;
+  kScales: MlxArray;
+  kZeros: MlxArray;
+  vPacked: MlxArray;
+  vScales: MlxArray;
+}
+
+export const disposeTurboQuant = (t: TurboQuantTensor): void => {
+  t.kIdx.dispose();
+  t.kScales.dispose();
+  t.kZeros.dispose();
+  t.vPacked.dispose();
+  t.vScales.dispose();
+};
+
+/** Numerical representation only. Storage owns capacity, positions, row
+ * membership and persistence; this codec retains the existing kernel choices. */
+export class TurboQuantCodec implements KvCodec<TurboQuantTensor> {
+  constructor(readonly kBits: number, readonly vBits: number, readonly fusedDecode: boolean) {}
+
+  /** Encode k/v [B,H,L,D] into a fresh TurboQuantTensor (packed at rest). */
+  encode(k: MlxArray, v: MlxArray): TurboQuantTensor {
+    const signed = this.kBits === 8;
+    const kEnc = tq.encodeKeys(k, this.kBits, signed);
+    const kIdxPacked = this.kBits < 8 ? tq.packBits(kEnc.indices, this.kBits) : kEnc.indices;
+    if (kIdxPacked !== kEnc.indices) kEnc.indices.dispose();
+    const vEnc = tq.encodeValues(v, this.vBits);
+    const vIdxPacked = this.vBits < 8 ? tq.packBits(vEnc.indices, this.vBits) : vEnc.indices;
+    if (vIdxPacked !== vEnc.indices) vEnc.indices.dispose();
+    return {
+      kIdx: kIdxPacked, kScales: kEnc.scales, kZeros: kEnc.zeros,
+      vPacked: vIdxPacked, vScales: vEnc.scales,
+    };
+  }
+
+  /** Decode a TurboQuantTensor's active [B,H,upTo,*] window back to
+   *  bf16 [B,H,upTo,headDim] (k, v). Caller owns the returned arrays.
+   *  deferV leaves V in the rotated domain (decodeValuesRotated) — the
+   *  caller un-rotates its attention output instead (tq.unrotateValues). */
+  decode(t: TurboQuantTensor, upTo: number, headDim: number, deferV = false): [MlxArray, MlxArray] {
+    const cut = (a: MlxArray): MlxArray => {
+      const [B, H, , D] = a.shape as [number, number, number, number];
+      return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
+    };
+    if (this.fusedDecode) {
+      const inputs = [cut(t.kIdx), cut(t.kScales), cut(t.kZeros), cut(t.vPacked), cut(t.vScales)] as const;
+      try {
+        const decoded = tq.tryDecodePackedKv(inputs, this.kBits, this.vBits, headDim, deferV);
+        if (decoded) return decoded;
+      } finally { for (const input of inputs) input.dispose(); }
+    }
+    const kIdxCut = cut(t.kIdx);
+    const kIdxUnpacked = this.kBits < 8 ? tq.unpackBits(kIdxCut, this.kBits, headDim) : kIdxCut;
+    if (kIdxUnpacked !== kIdxCut) kIdxCut.dispose();
+    const kScalesCut = cut(t.kScales);
+    const kZerosCut = cut(t.kZeros);
+    const k = tq.decodeKeys(kIdxUnpacked, kScalesCut, kZerosCut);
+    for (const a of [kIdxUnpacked, kScalesCut, kZerosCut]) a.dispose();
+
+    const vPackedCut = cut(t.vPacked);
+    const vIdxUnpacked = this.vBits < 8 ? tq.unpackBits(vPackedCut, this.vBits, headDim) : vPackedCut;
+    if (vIdxUnpacked !== vPackedCut) vPackedCut.dispose();
+    const vScalesCut = cut(t.vScales);
+    const v = deferV
+      ? tq.decodeValuesRotated(vIdxUnpacked, vScalesCut, this.vBits)
+      : tq.decodeValues(vIdxUnpacked, vScalesCut, this.vBits);
+    for (const a of [vIdxUnpacked, vScalesCut]) a.dispose();
+
+    return [k, v];
+  }
+
+}

package/src/model/universal/dense.ts

@@ -1,8 +1,10 @@
+import { captureKvAttention } from "../kv-attention-view";
 // UniversalDense — the Tier-0 config-driven dense llama-family module
 // (docs/design/generic-model-support.md §3.1). One module, shaped exactly
 // like qwen3.ts, with every arch delta selected by the UniversalArgs
 // descriptor (archs.ts). Monolith path only: no compiled
-// decode, no kv-quant — slow, never broken. The bar is L1 bit-exactness
+// decode. Standard attention uses the shared plain/affine KV port; the
+// softcap branch retains plain KV. The bar is L1 bit-exactness
 // vs mlx-lm on this machine's GPU (tests/parity/universal-parity.test.ts).
 //
 // Porting discipline: each branch transcribes its mlx-lm source op-for-op
@@ -154,19 +156,18 @@ class UniversalAttention {
       ? this.rope.applyDynamic(k, offArr)
       : this.rope.apply(k, cache.offset));
 
-    const [keys, values] = cache.updateAndFetch(k, v);
-    k.dispose();
-    v.dispose();
-
     let attn: MlxArray;
-    if (a.attnLogitSoftcap !== null) {
-      attn = this.#softcapAttention(q, keys, values, mask, B, L);
-    } else {
-      attn = ops.sdpa(q, keys, values, a.attnScale, mask.mode, mask.arr);
-    }
-    keys.dispose();
-    values.dispose();
-    q.dispose();
+    try {
+      if (a.attnLogitSoftcap !== null) {
+        const [keys, values] = cache.updateAndFetch(k, v);
+        try { attn = this.#softcapAttention(q, keys, values, mask, B, L); }
+        finally { keys.dispose(); values.dispose(); }
+      } else {
+        const view = captureKvAttention(cache, k, v);
+        try { attn = view.attend(q, a.attnScale, mask); }
+        finally { view.dispose(); }
+      }
+    } finally { k.dispose(); v.dispose(); q.dispose(); }
 
     const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
     attn.dispose();