mlx-bun 0.2.0 → 0.4.0

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

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

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/support.ts

@@ -56,7 +56,7 @@ export function isDrafterModelType(modelType: string): boolean {
 
 /** DiffusionGemma (model_type `diffusion_gemma`): block/masked-diffusion canvas
  *  model. Non-autoregressive — routed through the diffusion engine, not the AR
- *  loop. See docs/design/diffusion-gemma-port.md. */
+ *  loop. See docs/design/generic-model-support.md. */
 export function isDiffusionGemmaConfig(config: ModelConfig): boolean {
   return config.modelType === "diffusion_gemma";
 }

package/src/model/trellis-balanced-scatter.ts

@@ -0,0 +1,51 @@
+// Experimental variant 8, k=3 / T=256 / L<=12 only.
+// The word-per-lane kernel uses 24 of 32 SIMD lanes for a 3-bit block.
+// Assign eight outputs to every lane, load adjacent packed words, and align
+// them once so each unrolled decoder has constant shifts. Row accumulation,
+// split boundaries and the separate partial reduction stay identical to v6.
+export const BALANCED_SCATTER_SOURCE = String.raw`
+  const uint lane = thread_index_in_simdgroup;
+  const uint sg = simdgroup_index_in_threadgroup;
+  const uint block = thread_position_in_grid.y * (uint)SG_TG + sg;
+  const uint nBlocks = (uint)C / (uint)BT;
+  const uint zi = thread_position_in_grid.z;
+  const uint sample = zi / (uint)SPLITS;
+  const uint split = zi - sample * (uint)SPLITS;
+  if (block >= nBlocks || sample >= (uint)M) return;
+  const uint wpb = (uint)(BT * K / 32);
+  const ulong rowWords = (ulong)(INTERLEAVE ? 2 : nBlocks) * wpb;
+  const uint pBase = lane * 8u * (uint)K;
+  const uint wi = pBase >> 5u;
+  const uint off = pBase & 31u;
+  const uint next = (wi + 1u == wpb) ? 0u : wi + 1u;
+  const device uint32_t* words = codes + (INTERLEAVE
+    ? (ulong)(block / 2) * R * 2 * wpb + (block % 2) * wpb : (ulong)block * wpb);
+  const uint rowsPer = ((uint)R + (uint)SPLITS - 1u) / (uint)SPLITS;
+  const uint r0 = split * rowsPer;
+  const uint r1 = metal::min((uint)R, r0 + rowsPer);
+  const device T* xs = x + (ulong)sample * (ulong)R;
+  float acc[8];
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < 8u; ++i) acc[i] = 0.0f;
+  for (uint r = r0; r < r1; ++r) {
+    const uint word = words[(ulong)r * rowWords + wi];
+    const uint nxt = words[(ulong)r * rowWords + next];
+    const uint lo = off == 0u ? word : (word >> off) | (nxt << (32u - off));
+    const uint hi = off == 0u ? nxt : nxt >> off;
+    const float xr = float(xs[r]);
+    const float sr = float(scales[r]);
+    #pragma clang loop unroll(full)
+    for (uint i = 0; i < 8u; ++i) {
+      const uint shift = i * (uint)K;
+      uint win = lo >> shift;
+      if (shift + (uint)L > 32u) win |= hi << (32u - shift);
+      win &= (1u << (uint)L) - 1u;
+      const float value = (float)trellis_y(win) * (1.0f / 147.800537109375f);
+      acc[i] = metal::fma(value * sr, xr, acc[i]);
+    }
+  }
+  const ulong outBase = ((ulong)sample * (ulong)SPLITS + split) * (ulong)C;
+  #pragma clang loop unroll(full)
+  for (uint i = 0; i < 8u; ++i)
+    partial[outBase + block * (uint)BT + (uint)BT - 1u - lane * 8u - i] = acc[i];
+`;