mlx-bun 0.3.0 → 0.4.0

package/src/model/delayed-rotating-quantized-kv.ts

@@ -0,0 +1,93 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { RotatingKVCache, type Cache, type KvAttentionState, type KvAttentionView, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { SpeculativeTransitioningKvRows } from "../backends/mlx/transitioning-kv-rows";
+import { AlignedRotatingCache, alignRotatingRows, RotatingAffineLayout, SpeculativeRotatingAffineLayout, RotatingKvPositions } from "./rotating-kv-layout";
+import { BatchedRotatingQuantCache } from "./batched-rotating-quant";
+import { captureKvAttention, combineKvDonorAttention } from "./kv-attention-view";
+
+/** Precision changes preserve each row's physical columns. The model owns
+ * queries and scale; the shared lifecycle owns membership and conversion. */
+export class DelayedRotatingQuantizedKVCache extends SpeculativeTransitioningKvRows<RotatingAffineLayout> implements KvAttentionState, PaddedPrefillCache {
+  constructor(readonly maxSize: number, readonly groupSize: number, readonly bits: number, readonly start: number,
+    readonly maintain: (rows: Cache[]) => void, row?: Cache, readonly speculative = false) {
+    super({ signature: `kv:delayed-rotating-quant:${maxSize}:${bits}:${groupSize}:${start}`, maintain,
+      converted: row => row instanceof AlignedRotatingCache && row.inner instanceof BatchedRotatingQuantCache,
+      makeLayout: () => speculative ? new SpeculativeRotatingAffineLayout(maxSize, groupSize, bits) : new RotatingAffineLayout(maxSize, groupSize, bits), prepareRows: alignRotatingRows,
+      packRows: (layout, rows) => layout.adoptAlignedRows(rows as readonly AlignedRotatingCache[]),
+      extractRow: row => (row as AlignedRotatingCache).extract(speculative ? maxSize : undefined),
+      rollbackRow: (row, before, keep, preserve) => (row as AlignedRotatingCache).rollback(before, keep, preserve) }, row, new RotatingKvPositions(maxSize));
+  }
+  get attentionState(): KvAttentionState { return this; }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    if (this.packed) return this.packed.captureDonorAttention();
+    const views: import("./gemma4-base").KvDonorAttention[] = [];
+    try {
+      for (const row of this.rows) views.push((row as AlignedRotatingCache).captureDonorAttention());
+      return combineKvDonorAttention(views);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+  }
+  override makeMask(tokens: number, window: number | null) {
+    if (this.speculative && this.leftPad.every(pad => pad === 0) &&
+        this.offset + tokens <= Math.min(this.maxSize, window ?? this.maxSize))
+      return { mode: tokens === 1 ? "" as const : "causal" as const, arr: null };
+    return super.makeMask(tokens, window);
+  }
+  preparePrefill(padding: PrefillPadding): void {
+    if (!this.batchSize) {
+      const rows = padding.lengths.map(() => new RotatingKVCache(this.maxSize));
+      try { this.mergeRows(rows); } finally { for (const row of rows) row.dispose(); }
+    }
+    if (this.packed) { this.packed.preparePrefill(padding); return; }
+    // Every physical row follows the same block geometry, even when only a
+    // sibling has right padding. The token lengths still belong to each row.
+    const rightPaddedBatch = padding.rightPadding?.some(pad => pad > 0);
+    for (const [row, cache] of this.rows.entries()) (cache as AlignedRotatingCache).inner.preparePrefill({
+      lengths: [padding.lengths[row]!],
+      ...(padding.leftPadding ? { leftPadding: [padding.leftPadding[row]!] } : {}),
+      ...(padding.rightPadding ? { rightPadding: [padding.rightPadding[row]!] } : {}),
+    }, rightPaddedBatch);
+    this.syncPositions(true);
+  }
+  finalizePrefill(): void {
+    if (this.packed) { this.packed.finalizePrefill(); return; }
+    for (const cache of this.rows) (cache as AlignedRotatingCache).inner.finalizePrefill();
+    this.syncPositions(true);
+  }
+  get maxTokens(): number { return this.maxSize; }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
+  makeEmptyBatch(): DelayedRotatingQuantizedKVCache {
+    return new DelayedRotatingQuantizedKVCache(this.maxSize, this.groupSize, this.bits, this.start, this.maintain, undefined, this.speculative);
+  }
+  updateAndFetch(): [MlxArray, MlxArray] { throw new Error("mixed precision rows use their attention state"); }
+  releaseRopeArr(): void { this.packed?.releaseRopeArr(); }
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
+    this.advance();
+    if (this.packed) return captureKvAttention(this.packed, k, v);
+    const views: KvAttentionView[] = [];
+    try {
+      for (const [row, cache] of this.rows.entries()) {
+        const slice = (a: MlxArray) => a.slice([row, 0, 0, 0], [row + 1, a.shape[1]!, a.shape[2]!, a.shape[3]!]);
+        using kr = slice(k), vr = slice(v);
+        views.push(captureKvAttention(cache, kr, vr));
+      }
+      this.syncPositions(true);
+    } catch (error) { for (const view of views) view.dispose(); throw error; }
+    return {
+      attend(q, scale, mask) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, view] of views.entries()) {
+            using qr = q.slice([row, 0, 0, 0], [row + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+            const start = mask.arr?.shape.map(() => 0), end = mask.arr ? [...mask.arr.shape] : undefined;
+            if (start && end && end.length === 4 && end[0]! > 1) { start[0] = row; end[0] = row + 1; }
+            using arr = mask.arr ? mask.arr.slice(start!, end!) : null;
+            outputs.push(view.attend(qr, scale, { mode: mask.mode, arr }));
+          }
+          return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+        } finally { for (const output of outputs) output.dispose(); }
+      },
+      dispose() { for (const view of views) view.dispose(); },
+    };
+  }
+}

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

@@ -0,0 +1,55 @@
+import { captureFullKvDonorRows } from "../backends/mlx/full-kv-row-donor";
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import { appendFullKvRows } from "../backends/mlx/full-kv-row-append";
+import type { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { unrotateValues } from "../mlx/turboquant-ops";
+import { type Cache, type RotatedValueAttentionState, TurboQuantKVCache } from "./gemma4-base";
+import { BatchedTurboQuantKVCache } from "./batched-turboquant-kv";
+import { FullTransitioningKvRows } from "../backends/mlx/full-transitioning-kv-rows";
+
+/** TurboQuant tensor access and value-domain handling over the shared
+ * row-transition lifecycle. Model projections and attention remain batched. */
+export class DelayedTurboQuantKVCache extends FullTransitioningKvRows<BatchedTurboQuantKVCache> implements RotatedValueAttentionState {
+  #rotated: boolean[] = [];
+  constructor(readonly kBits: number, readonly vBits: number, readonly start: number,
+    readonly maintain: (rows: Cache[]) => void, row?: Cache) {
+    super({ signature: `kv:delayed-turboquant:${kBits}:${vBits}:${start}`, maintain,
+      converted: row => row instanceof TurboQuantKVCache,
+      makeLayout: () => new BatchedTurboQuantKVCache(kBits, vBits) }, row);
+  }
+  captureDonorRows(): import("./gemma4-base").KvDonorRows {
+    return this.packed?.captureDonorRows() ?? captureFullKvDonorRows(this.rows, this.leftPad, this.offset);
+  }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    return decodedKvDonorAttention(this.captureDonorRows());
+  }
+  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
+  #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 => !!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.maintain); }
+}

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

@@ -0,0 +1,56 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { PrefillPadding } from "./gemma4-base";
+import { rowRollIndices } from "./batched-row-storage";
+import { disposeResources } from "../engine/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/model/full-prefill-row.ts

@@ -0,0 +1,105 @@
+import type { MlxArray } from "../mlx/array";
+import { materializeCopy } from "../mlx/materialize";
+import { withResource, disposeResources } from "../engine/resources";
+import { leaseCacheState } from "../backends/mlx/state-views";
+import { FullPrefillPadding } from "./full-prefill-padding";
+import { KVCache, QuantizedKVCache, TurboQuantKVCache, type Cache, type Mask, type PrefillPadding } from "./gemma4-base";
+
+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 (this.inner as Cache).quantizedAttention; }
+  get rotatedValueAttention() { return (this.inner as Cache).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();
+      if (planes.length) result.restoreState(planes[0]!, planes[1]!, offset);
+      return result;
+    }
+    if (source instanceof QuantizedKVCache) {
+      const result = new QuantizedKVCache(source.groupSize, source.bits);
+      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);
+    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); }
+  }
+  toQuantized(groupSize: number, bits: number): FullPrefillRow {
+    this.inner = (this.inner as KVCache).toQuantized(groupSize, bits);
+    return this;
+  }
+  toTurboQuantized(kBits: number, vBits: number): FullPrefillRow {
+    this.inner = TurboQuantKVCache.fromKVCache(this.inner as KVCache, kBits, vBits);
+    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; }