mlx-bun 0.3.0 → 0.5.0

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

@@ -0,0 +1,94 @@
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
+import type { MlxArray } from "../mlx/array";
+import { KvTensorRows, type KvTensorRowView } from "../backends/mlx/kv-tensor-rows";
+import { disposeResources } from "../engine/resources";
+import { TurboQuantKVCache, type BatchableCache, type Cache, type Mask, type RotatedValueAttentionState, type PaddedPrefillCache, type PrefillPadding } from "./gemma4-base";
+import { TurboQuantCodec, turboQuantFusedDecode, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
+
+const fields = ["kIdx", "kScales", "kZeros", "vPacked", "vScales"] as const;
+function encoded(planes: readonly MlxArray[]): TurboQuantTensor {
+  return { kIdx: planes[0]!, kScales: planes[1]!, kZeros: planes[2]!, vPacked: planes[3]!, vScales: planes[4]! };
+}
+
+/** The existing TurboQuant codec over shared tensor-row storage. Row changes
+ * copy encoded bytes without decoding, rotating or requantizing their values. */
+export class BatchedTurboQuantKVCache implements BatchableCache, PaddedPrefillCache {
+  readonly #storage = new KvTensorRows();
+  readonly #codec: TurboQuantCodec;
+  #reuseOffsets: number[] = [];
+  #headDim: number | null = null;
+  constructor(readonly kBits: number, readonly vBits: number,
+    readonly fusedDecode = turboQuantFusedDecode()) {
+    this.#codec = new TurboQuantCodec(kBits, vBits, fusedDecode);
+  }
+  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#storage.restorePrefillEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    this.#storage.preparePrefill(padding);
+    if (!this.#reuseOffsets.length) this.#reuseOffsets = padding.lengths.map(() => 0);
+  }
+  finalizePrefill(): void { this.#storage.finalizePrefill(); }
+  captureDonorRows(): import("./gemma4-base").KvDonorRows {
+    if (!this.#storage.planes.length || this.#headDim === null) throw new Error("cache is empty");
+    const [keys, values] = this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, false);
+    return { keys, values, ...this.#storage.captureDonorValidity() };
+  }
+  captureDonorAttention(): import("./gemma4-base").KvDonorAttention {
+    return decodedKvDonorAttention(this.captureDonorRows());
+  }
+  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
+  signature(): string { return `kv:batched-turboquant:${this.kBits}:${this.vBits}`; }
+  get minimumReusableOffset(): number { return Math.max(0, ...this.#reuseOffsets); }
+  get headDim(): number | null { return this.#headDim; }
+  get rowOffsets(): readonly number[] { return this.#storage.rowOffsets; }
+  get leftPad(): readonly number[] { return this.#storage.leftPad; }
+  get batchSize(): number | null { return this.#storage.batchSize; }
+  get offset(): number { return this.#storage.offset; }
+  get ropeOffsetArr(): MlxArray | undefined { return this.#storage.ropeOffsetArr; }
+  makeEmptyBatch(): BatchedTurboQuantKVCache { return new BatchedTurboQuantKVCache(this.kBits, this.vBits, this.fusedDecode); }
+  bytesPerToken(): number { return this.#storage.bytesPerToken(); }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+  makeMask(tokens: number, window: number | null): Mask { return this.#storage.makeMask(tokens, window); }
+  isTrimmable(): boolean { return true; }
+  trim(count: number): void { this.#storage.trim(count); }
+  specRoundBegin(): void { this.#storage.specRoundBegin(); }
+  specRoundCommit(): void { this.#storage.specRoundCommit(); }
+  specRoundRollback(keep: number | readonly number[]): void { this.#storage.specRoundRollback(keep); }
+  #append(k: MlxArray, v: MlxArray, defer: boolean): [MlxArray, MlxArray] {
+    const packed = this.#codec.encode(k, v);
+    try {
+      this.#storage.append(fields.map(field => packed[field]));
+      this.#headDim = k.shape[3]!;
+      return this.#codec.decode(encoded(this.#storage.planes), this.offset, this.#headDim, defer);
+    } finally { disposeTurboQuant(packed); }
+  }
+  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); }
+  mergeRows(rows: readonly Cache[]): void {
+    const views: KvTensorRowView[] = [], held: MlxArray[] = [];
+    const reuseOffsets: number[] = [];
+    let dim = this.#headDim;
+    try {
+      for (const row of rows) {
+        if (row instanceof BatchedTurboQuantKVCache) {
+          views.push(row.#storage); dim ??= row.headDim; reuseOffsets.push(...row.#reuseOffsets);
+        } else if (row instanceof TurboQuantKVCache) {
+          const planes = row.state(); held.push(...planes); reuseOffsets.push(row.minimumReusableOffset);
+          views.push({ planes, rowOffsets: [row.offset], leftPad: [0] }); dim ??= row.headDim;
+        } else throw new Error(`TurboQuant row layout cannot merge ${row.signature()}`);
+      }
+      this.#storage.mergeRows(views); this.#headDim = dim; this.#reuseOffsets = reuseOffsets;
+    } finally { disposeResources(held); }
+  }
+  alignRows(leftPad: readonly number[]): void { this.#storage.alignRows(leftPad); }
+  filterRows(keep: readonly number[]): void {
+    this.#storage.filterRows(keep); this.#reuseOffsets = keep.map(row => this.#reuseOffsets[row]!);
+  }
+  state(): MlxArray[] { return this.#storage.planes; }
+  extractRow(row: number): TurboQuantKVCache {
+    const result = new TurboQuantKVCache(this.kBits, this.vBits, this.fusedDecode), count = this.rowOffsets[row]!;
+    if (count) result.restoreState(encoded(this.#storage.extractRow(row)), count, this.#headDim!);
+    result.minimumReusableOffset = this.#reuseOffsets[row] ?? 0;
+    return result;
+  }
+  dispose(): void { this.#storage.dispose(); this.#headDim = null; this.#reuseOffsets = []; }
+}

package/src/model/compiled-decode.ts

@@ -277,7 +277,10 @@ function decodePhase(c: AnyCache): "concat" | "ring" {
 }
 
 function disposeShared(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {
@@ -381,6 +384,9 @@ export class CompiledDecode {
   readonly #segmented: boolean;
 
   private constructor(readonly model: Gemma4Model) {
+    // A fresh resume has no prefill to evaluate lazy FP32 constants. Capture
+    // their values, not their construction graph, in every compiled closure.
+    model.materializeGraphConstants();
     const t = model.config.text;
     this.#segmented =
       t.numKvSharedLayers === 0 && !t.enableMoeBlock && model.perLayerWidth === 0;

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

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

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

@@ -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}`, 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

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

package/src/model/factory.ts

@@ -23,6 +23,7 @@ import {
   type Glm52MemoryPlan,
 } from "./glm52-memory";
 import { UniversalDenseModel } from "./universal/dense";
+import { WhisperModel } from "./whisper";
 import { genericArgsFor } from "./universal/archs";
 import {
   ModelImplementationRegistry,
@@ -134,6 +135,26 @@ export async function openGlm52RuntimeModel(
   return { model, plan };
 }
 
+/** Open a Whisper checkpoint (encoder-decoder speech model) through the same
+ * profile resolution as text models. Whisper never enters the chat loop, so
+ * it is not a RuntimeModel: the transcription engine (src/audio/whisper-*)
+ * owns it. */
+export async function openWhisperModel(
+  modelDir: string, options: { readonly profiles?: ResolveModelProfileOptions } = {},
+): Promise<{ model: WhisperModel; profile: ResolvedModelProfile; config: ModelConfig }> {
+  const config = await loadModelConfig(modelDir);
+  const profile = resolveModelProfile(config, options.profiles);
+  if (profile.profile.execution.graph !== "whisper")
+    throw new Error(`${modelDir} is not a Whisper checkpoint (profile ${profile.profile.id})`);
+  const weights = await Weights.open(modelDir);
+  try {
+    return { model: new WhisperModel(weights, config), profile, config };
+  } catch (error) {
+    weights.dispose();
+    throw error;
+  }
+}
+
 /**
  * Artifact-aware construction. GLM-5.2's Colibri snapshot has no ordinary
  * model.safetensors.index.json, so it must bypass Weights.open and use its
@@ -155,6 +176,11 @@ export async function openModel<Model>(
       throw new Error("custom Colibri implementations require a streamed loader binding; refusing to fall back");
     return (await openGlm52RuntimeModel(modelDir, options)).model;
   }
+  if (profile.profile.execution.loop === "encoder-decoder")
+    throw new Error(
+      `${modelDir} is a speech model (${profile.profile.id}); open it with openWhisperModel ` +
+      "or serve it as a transcription model — it has no chat/completion loop",
+    );
   // Resolve before opening weights; missing or incompatible code allocates nothing.
   const implementation = (options.implementations ?? MLX_MODEL_IMPLEMENTATIONS).select(config, profile);
   const weights = await Weights.open(modelDir);

package/src/model/fingerprint.ts

@@ -31,3 +31,15 @@ export function configFingerprint(config: ModelConfig): string {
   h.update(payload);
   return h.digest("hex").slice(0, 16);
 }
+
+/** Exact immutable model-component identity, calculated once while its owner
+ * holds the native execution lease. Names/shapes/dtypes disambiguate layouts. */
+export function tensorFingerprint(entries: Iterable<readonly [string,
+  Pick<import("../mlx/array").MlxArray, "dtype" | "shape" | "rawBytesView">]>): string {
+  const hash = new Bun.CryptoHasher("sha256");
+  for (const [name, tensor] of [...entries].sort(([a], [b]) => a.localeCompare(b))) {
+    hash.update(JSON.stringify([name, tensor.dtype, tensor.shape]));
+    hash.update(tensor.rawBytesView());
+  }
+  return hash.digest("hex");
+}

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, source.fusedDecode);
+    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; }