mlx-bun 0.5.0 → 1.0.0

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

@@ -1,12 +0,0 @@
-import * as ops from "../mlx/ops";
-import type { KvDonorAttention, KvDonorRows } from "./gemma4-base";
-
-/** Takes ownership of decoded donor planes. Assistant attention uses decoded
- * values, preserving its SDPA arithmetic rather than moving inverse rotation. */
-export function decodedKvDonorAttention(view: KvDonorRows): KvDonorAttention {
-  const { keys, values, offsets, starts, ends } = view;
-  return { width: keys.shape[2]!, dtype: keys.dtype, offsets, starts, ends,
-    attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
-    dispose() { keys.dispose(); values.dispose(); },
-  };
-}

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

@@ -1,72 +0,0 @@
-import { captureFullKvDonorAttention } from "../backends/mlx/full-kv-row-donor";
-import { appendFullKvRows } from "../backends/mlx/full-kv-row-append";
-import { fullRowPadding, fullRowPhysicalLength } from "./full-prefill-row";
-import { captureKvAttention } from "./kv-attention-view";
-import type { MlxArray } from "../mlx/array";
-import * as ops from "../mlx/ops";
-import { FullTransitioningKvRows } from "../backends/mlx/full-transitioning-kv-rows";
-import { BatchedQuantizedKVCache } from "./batched-quantized-kv";
-import { QuantizedKVCache, type Cache, type Mask, type KvAttentionState, type KvAttentionView } from "./gemma4-base";
-
-/** Affine storage keeps its native quantized-attention arithmetic while rows
- * cross the conversion boundary independently. Once all rows convert, their
- * packed planes use the existing batched attention implementation. */
-export class DelayedQuantizedKVCache extends FullTransitioningKvRows<BatchedQuantizedKVCache> implements KvAttentionState {
-  constructor(readonly groupSize: number, readonly bits: number, readonly start: number,
-    readonly maintain: (rows: Cache[]) => void, row?: Cache) {
-    super({ signature: `kv:delayed-quant:${bits}:${groupSize}:${start}`, conversionOffset: start, maintain,
-      converted: row => row instanceof QuantizedKVCache,
-      makeLayout: () => new BatchedQuantizedKVCache(groupSize, bits) }, row);
-  }
-  get attentionState(): KvAttentionState { return this; }
-  captureDonorAttention() {
-    return this.packed?.captureDonorAttention() ?? captureFullKvDonorAttention(this.rows, this.leftPad, this.offset);
-  }
-  makeEmptyBatch(): DelayedQuantizedKVCache { return new DelayedQuantizedKVCache(this.groupSize, this.bits, this.start, this.maintain); }
-  updateAndFetch(): [MlxArray, MlxArray] { throw new Error("mixed precision rows use their attention state"); }
-  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
-    this.advance();
-    if (this.packed) return captureKvAttention(this.packed, k, v);
-    if (this.rows.every(row => !row.quantizedAttention)) {
-      const [keys, values] = appendFullKvRows(this.rows, this.leftPad, this.offset + k.shape[2]!, k, v);
-      try { this.syncPositions(true); } catch (error) { keys.dispose(); values.dispose(); throw error; }
-      return { attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
-        dispose() { keys.dispose(); values.dispose(); } };
-    }
-    const views: KvAttentionView[] = [], pads = this.rows.map((row, index) => this.leftPad[index]! - fullRowPadding(row));
-    const lengths = this.rows.map(cache => fullRowPhysicalLength(cache) + k.shape[2]!);
-    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, independentPositions) {
-        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]!]);
-            // Capture alignment with the KV tensors, never consult live rows.
-            const start = mask.arr?.shape.map(() => 0), end = mask.arr ? [...mask.arr.shape] : undefined;
-            if (start && end) {
-              if (end.length === 4 && end[0]! > 1) { start[0] = row; end[0] = row + 1; }
-              start[start.length - 1] = pads[row]!;
-              end[end.length - 1] = pads[row]! + lengths[row]!;
-            }
-            using arr = mask.arr ? mask.arr.slice(start!, end!) : null;
-            outputs.push(view.attend(qr, scale, { mode: mask.mode, arr }, independentPositions));
-          }
-          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(); },
-    };
-  }
-  updateAndAttend(q: MlxArray, k: MlxArray, v: MlxArray, scale: number, mask: Mask): MlxArray {
-    const view = this.appendAndFetch(k, v);
-    try { return view.attend(q, scale, mask); } finally { view.dispose(); }
-  }
-}

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

@@ -1,93 +0,0 @@
-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}`, conversionOffset: 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

@@ -1,57 +0,0 @@
-import { turboQuantFusedDecode } from "./turboquant-codec";
-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,
-    readonly fusedDecode = turboQuantFusedDecode()) {
-    super({ signature: `kv:delayed-turboquant:${kBits}:${vBits}:${start}`, conversionOffset: start, maintain,
-      converted: row => row instanceof TurboQuantKVCache,
-      makeLayout: () => new BatchedTurboQuantKVCache(kBits, vBits, fusedDecode) }, 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, undefined, this.fusedDecode); }
-}