mlx-bun 0.4.0 → 1.0.0

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

@@ -1,56 +0,0 @@
-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

@@ -1,105 +0,0 @@
-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; }