mlx-bun 0.4.0 → 0.5.0

package/src/mlx/ops.ts

@@ -472,14 +472,12 @@ export function where(cond: MlxArray, x: MlxArray, y: MlxArray, s: S = gpuStream
   );
 }
 
-// arange is built host-side and uploaded. The original mlx_arange binding
-// broke once the calling path got JIT-optimized; a standalone repro
-// (lab/repro/bun-ffi-f64/) traced the root cause to a Bun bug: after DFG
-// tier-up, typed-array reads following a bun:ffi call return stale values
-// (the JIT eliminates the load across the native call). Not f64 marshaling
-// — args reach C intact. See PLAN.md Phase 4 findings. Reuse large ranges
-// across attention layers and sampling, but bound their retained storage:
-// causal masks request a new range on every round once context exceeds 64K.
+// The old native arange workaround concerned stale typed-array out-pointer
+// reads after Bun JIT tier-up, not f64 argument marshaling. outArray uses
+// native pointer reads now. Large nonnegative unit-step int32 ranges avoid
+// host construction; small/general integer ranges retain JS conversion,
+// including its wrapping behavior outside int32. Cache storage stays bounded
+// as causal masks request new lengths on successive long-context rounds.
 const arangeCache = new Map<string, { array: MlxArray; bytes: number }>();
 const ARANGE_CACHE_MIN = 65536;
 const ARANGE_CACHE_BYTES = 8 * 1024 * 1024;
@@ -498,8 +496,10 @@ export function arange(start: number, stop: number, step: number, dtype: Dtype,
   }
 
   let arr: MlxArray;
-  if (dtype === Dtype.int32) {
-    // host-built int32 leaf (the Bun-f64-FFI-safe path — see note above).
+  const nativePositions = n === 0 ||
+    (n >= ARANGE_CACHE_MIN && step === 1 && start >= 0 && stop <= 2147483647);
+  if (dtype === Dtype.int32 && !nativePositions) {
+    // Preserve JS conversion for small/general integer ranges.
     const data = new Int32Array(n);
     for (let i = 0; i < n; i++) data[i] = start + i * step;
     arr = MlxArray.fromInt32(data, [n]);
@@ -1019,3 +1019,90 @@ export function leftShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray
 export function rightShift(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
   return new MlxArray(outArray("right_shift", (o) => C.mlx_right_shift(o, a.handle, b.handle, s)));
 }
+
+// --- Whisper additions (fft / strided framing / pad / reductions) ----------
+
+/** mx.fft.rfft(a, n, axis) with numpy "backward" normalization. Real input
+ *  → complex64 output of length n/2+1 along `axis`. */
+export function rfft(a: MlxArray, n: number, axis = -1, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("fft_rfft", (o) => C.mlx_fft_rfft(o, a.handle, n, axis, 0, s)));
+}
+
+/** mx.as_strided(a, shape, strides, offset) — element strides. */
+export function asStrided(
+  a: MlxArray, shape: number[], strides: number[], offset = 0, s: S = gpuStream,
+): MlxArray {
+  const shp = new Int32Array(shape);
+  const str = new BigInt64Array(strides.map((v) => BigInt(v)));
+  return new MlxArray(
+    outArray("as_strided", (o) =>
+      C.mlx_as_strided(o, a.handle, ptr(shp), BigInt(shape.length), ptr(str), BigInt(strides.length), BigInt(offset), s),
+    ),
+  );
+}
+
+/** mx.pad(a, [(lo, hi), …]) with constant zero fill (all axes given). */
+export function pad(a: MlxArray, widths: [number, number][], s: S = gpuStream): MlxArray {
+  const axes = new Int32Array(widths.map((_, i) => i));
+  const lo = new Int32Array(widths.map((w) => w[0]));
+  const hi = new Int32Array(widths.map((w) => w[1]));
+  const zero = scalarLike(0, a);
+  try {
+    return new MlxArray(
+      outArray("pad", (o) =>
+        C.mlx_pad(
+          o, a.handle, ptr(axes), BigInt(widths.length), ptr(lo), BigInt(widths.length),
+          ptr(hi), BigInt(widths.length), zero.handle, ptr(cstr("constant")), s,
+        ),
+      ),
+    );
+  } finally {
+    zero.dispose();
+  }
+}
+
+export function log10(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("log10", (o) => C.mlx_log10(o, a.handle, s)));
+}
+
+/** Reduce-max over all elements → 0-d array (keepdims=false). */
+export function maxAll(a: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("max", (o) => C.mlx_max(o, a.handle, false, s)));
+}
+
+export function broadcastTo(a: MlxArray, shape: number[], s: S = gpuStream): MlxArray {
+  const buf = new Int32Array(shape);
+  return new MlxArray(
+    outArray("broadcast_to", (o) => C.mlx_broadcast_to(o, a.handle, ptr(buf), BigInt(shape.length), s)),
+  );
+}
+
+/** mx.topk(a, k, axis) — the k largest values along axis (unsorted). */
+export function topkAxis(a: MlxArray, k: number, axis: number, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("topk_axis", (o) => C.mlx_topk_axis(o, a.handle, k, axis, s)));
+}
+
+export function varAxis(a: MlxArray, axis: number, keepdims: boolean, ddof = 0, s: S = gpuStream): MlxArray {
+  return new MlxArray(
+    outArray("var_axis", (o) => C.mlx_var_axis(o, a.handle, axis, keepdims, ddof, s)),
+  );
+}
+
+export function stackAxis(arrays: MlxArray[], axis: number, s: S = gpuStream): MlxArray {
+  const handles = new BigUint64Array(arrays.map((a) => a.handle));
+  const vec = C.mlx_vector_array_new_data(ptr(handles), BigInt(arrays.length));
+  try {
+    return new MlxArray(outArray("stack_axis", (o) => C.mlx_stack_axis(o, vec, axis, s)));
+  } finally {
+    C.mlx_vector_array_free(vec);
+  }
+}
+
+export function minimum(a: MlxArray, b: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("minimum", (o) => C.mlx_minimum(o, a.handle, b.handle, s)));
+}
+
+/** Flat mx.take(a, indices) — gathers from the flattened array. */
+export function takeFlat(a: MlxArray, indices: MlxArray, s: S = gpuStream): MlxArray {
+  return new MlxArray(outArray("take", (o) => C.mlx_take(o, a.handle, indices.handle, s)));
+}

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

@@ -3,7 +3,7 @@ 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, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
+import { TurboQuantCodec, turboQuantFusedDecode, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
 
 const fields = ["kIdx", "kScales", "kZeros", "vPacked", "vScales"] as const;
 function encoded(planes: readonly MlxArray[]): TurboQuantTensor {
@@ -17,8 +17,9 @@ export class BatchedTurboQuantKVCache implements BatchableCache, PaddedPrefillCa
   readonly #codec: TurboQuantCodec;
   #reuseOffsets: number[] = [];
   #headDim: number | null = null;
-  constructor(readonly kBits: number, readonly vBits: number) {
-    this.#codec = new TurboQuantCodec(kBits, vBits, process.env.MLX_BUN_TURBOQUANT_FUSED_DECODE === "1");
+  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 {
@@ -43,7 +44,7 @@ export class BatchedTurboQuantKVCache implements BatchableCache, PaddedPrefillCa
   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); }
+  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); }
@@ -84,7 +85,7 @@ export class BatchedTurboQuantKVCache implements BatchableCache, PaddedPrefillCa
   }
   state(): MlxArray[] { return this.#storage.planes; }
   extractRow(row: number): TurboQuantKVCache {
-    const result = new TurboQuantKVCache(this.kBits, this.vBits), count = this.rowOffsets[row]!;
+    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;

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

@@ -14,7 +14,7 @@ import { QuantizedKVCache, type Cache, type Mask, type KvAttentionState, type Kv
 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}`, maintain,
+    super({ signature: `kv:delayed-quant:${bits}:${groupSize}:${start}`, conversionOffset: start, maintain,
       converted: row => row instanceof QuantizedKVCache,
       makeLayout: () => new BatchedQuantizedKVCache(groupSize, bits) }, row);
   }
@@ -44,7 +44,7 @@ export class DelayedQuantizedKVCache extends FullTransitioningKvRows<BatchedQuan
       this.syncPositions(true);
     } catch (error) { for (const view of views) view.dispose(); throw error; }
     return {
-      attend(q, scale, mask) {
+      attend(q, scale, mask, independentPositions) {
         const outputs: MlxArray[] = [];
         try {
           for (const [row, view] of views.entries()) {
@@ -57,7 +57,7 @@ export class DelayedQuantizedKVCache extends FullTransitioningKvRows<BatchedQuan
               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 }));
+            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(); }

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

@@ -11,7 +11,7 @@ import { captureKvAttention, combineKvDonorAttention } from "./kv-attention-view
 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,
+    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[]),

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

@@ -1,3 +1,4 @@
+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";
@@ -14,10 +15,11 @@ import { FullTransitioningKvRows } from "../backends/mlx/full-transitioning-kv-r
 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,
+    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) }, row);
+      makeLayout: () => new BatchedTurboQuantKVCache(kBits, vBits, fusedDecode) }, row);
   }
   captureDonorRows(): import("./gemma4-base").KvDonorRows {
     return this.packed?.captureDonorRows() ?? captureFullKvDonorRows(this.rows, this.leftPad, this.offset);
@@ -51,5 +53,5 @@ export class DelayedTurboQuantKVCache extends FullTransitioningKvRows<BatchedTur
       return ops.where(mask, restored, output);
     };
   }
-  makeEmptyBatch(): DelayedTurboQuantKVCache { return new DelayedTurboQuantKVCache(this.kBits, this.vBits, this.start, this.maintain); }
+  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-row.ts

@@ -65,7 +65,7 @@ export class FullPrefillRow implements Cache {
         { packed: planes[3]!, scales: planes[4]!, biases: planes[5]! }, offset);
       return result;
     }
-    const result = new TurboQuantKVCache(source.kBits, source.vBits);
+    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;

package/src/model/gemma4-base.ts

@@ -12,13 +12,13 @@ import type { ModelConfig } from "../config";
 import { quantFor } from "../config";
 import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
+import { Dtype, deviceArchitecture } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { quantizedMatmulRows } from "../mlx/quantized-rows";
 import { expertOffloadArray } from "../expert-offload";
 import * as tq from "../mlx/turboquant-ops";
 import type { KvCodec } from "../backends/mlx/kv-codec";
-import { TurboQuantCodec, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
+import { TurboQuantCodec, turboQuantFusedDecode, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
 export { disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
 import { runtimeValue } from "../runtime-config";
 
@@ -297,7 +297,8 @@ export interface RotatedValueAttentionState {
  * again. It owns its tensor handles independently of later cache membership or
  * precision changes. Query/mask inputs are borrowed; outputs are owned. */
 export interface KvAttentionView {
-  attend(q: MlxArray, scale: number, mask: Mask): MlxArray;
+  /** Qualified committed spans retain the single-position reduction order. */
+  attend(q: MlxArray, scale: number, mask: Mask, independentPositions?: boolean): MlxArray;
   dispose(): void;
 }
 
@@ -325,6 +326,8 @@ export interface KvDonorAttention extends KvAttentionView {
 }
 
 export interface Cache {
+  /** Maximum committed positions before this state changes precision. */
+  maxAppendTokens?(): number;
   captureDonorRows?(): KvDonorRows;
   captureDonorAttention?(): KvDonorAttention;
   /** A preparation method commits each request's own precision boundary.
@@ -1671,8 +1674,9 @@ export class TurboQuantKVCache implements Cache {
   /** Experimental operation selection, captured once for this cache. */
   readonly #codec: KvCodec<TurboQuantTensor>;
 
-  constructor(readonly kBits: number, readonly vBits: number) {
-    this.#codec = new TurboQuantCodec(kBits, vBits, process.env.MLX_BUN_TURBOQUANT_FUSED_DECODE === "1");
+  constructor(readonly kBits: number, readonly vBits: number,
+    readonly fusedDecode = turboQuantFusedDecode()) {
+    this.#codec = new TurboQuantCodec(kBits, vBits, fusedDecode);
   }
 
   captureDonorRows(): KvDonorRows {
@@ -1940,8 +1944,23 @@ export function quantizedSdpaUnfused(
     vT = expand(vq);
   }
 
-  let scores = ops.quantizedMatmulQT(queries, kT, true, groupSize, bits);
+  // M4 Pro's native key matvec retains exact arithmetic when three GQA
+  // heads share a batch. Restore score geometry before softmax/value work.
+  const groupHeads = B <= 2 && H === 24 && KV === 4 && L === 3 && D === 256 &&
+    N >= 8192 && groupSize === 64 && bits === 4 &&
+    (q.dtype === Dtype.bfloat16 || q.dtype === Dtype.float32) &&
+    deviceArchitecture() === "applegpu_g16s";
+  let keyQueries = queries;
+  if (groupHeads) {
+    keyQueries = ops.reshape(queries, [B, KV, 2, 9, D]);
+    owned.push(keyQueries);
+  }
+  let scores = ops.quantizedMatmulQT(keyQueries, kT, true, groupSize, bits);
   owned.push(scores);
+  if (groupHeads) {
+    scores = ops.reshape(scores, [B, KV, nRep, L, N]);
+    owned.push(scores);
+  }
 
   let maskArr: MlxArray | null = null;
   let ownsMask = false;

package/src/model/gemma4.ts

@@ -52,11 +52,12 @@ import {
   captureRopeOffsets,
   type LoraWeights,
 } from "./gemma4-base";
+import { mapPackedTokens, type TokenGroup } from "./token-groups";
 import { Checkpoint } from "../mlx/checkpoint";
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 import { unrotateValues as tqUnrotateValues } from "../mlx/turboquant-ops";
 import { CompiledFunction } from "../mlx/compile";
-import { flagOn } from "../runtime-config";
+import { runtimeFlag } from "../runtime-config";
 
 /** Verbatim port of mlx_lm/models/gemma4_text.py:
  *    `@partial(mx.compile, shapeless=True) def geglu(gate, x): return nn.gelu_approx(gate) * x`
@@ -77,7 +78,7 @@ import { flagOn } from "../runtime-config";
  *  ~9-op path. `MLX_BUN_COMPILED_GEGLU=0` (--compiled-activations off) selects the
  *  uncompiled composition — same L1 parity, slower (the A/B opt-out). */
 export function compiledGegluActive(): boolean {
-  return flagOn("MLX_BUN_COMPILED_GEGLU", true);
+  return runtimeFlag("MLX_BUN_COMPILED_GEGLU", true);
 }
 
 let _gegluClosure: CompiledFunction | null = null;
@@ -1020,6 +1021,85 @@ export class Gemma4Model {
     return h;
   }
 
+  /** Mixed token execution: attention retains each group's cache, positions and
+   * SDPA shape; feed-forward work packs only real tokens. Kernel selection for
+   * packed matmuls can differ from a solo GEMV, so this is an explicit Lab port
+   * until its matching packed oracle and serving comparison are accepted. */
+  forwardHiddenMixed(work: readonly TokenGroup[]): MlxArray[] {
+    if (work.length === 1 && !work[0]!.captureLayer) return [this.forwardHidden(work[0]!.ids, work[0]!.cache)];
+    const groups: Array<{ ids: MlxArray; cache: Cache[]; h: MlxArray;
+      masks: Map<string, Mask>; perLayer: MlxArray | null; intermediates: (SharedKv | null)[] }> = [];
+    const results: MlxArray[] = [];
+    const packedMlp = runtimeFlag("MLX_BUN_MIXED_PACKED_MLP", true) && !work.some(group => group.preserveTokenGeometry);
+    try {
+      for (const { ids, cache } of work) {
+        const group = { ids, cache, h: this.embed.encode(ids), masks: new Map<string, Mask>(),
+          perLayer: null as MlxArray | null, intermediates: Array<SharedKv | null>(this.layers.length).fill(null) };
+        groups.push(group);
+        group.h = disposing(group.h, ops.mulScalar(group.h, this.embedScale));
+        for (let i = 0; i < this.numDonors; i++) {
+          const type = this.layers[i]!.layerType;
+          if (!group.masks.has(type)) group.masks.set(type,
+            cache[i]!.makeMask(ids.shape[1]!, type === "sliding_attention" ? this.windowSize : null));
+        }
+        if (this.perLayerWidth > 0) group.perLayer = this.computePerLayerInputs(ids, group.h);
+      }
+      for (let i = 0; i < this.layers.length; i++) {
+        const layer = this.layers[i]!, ci = this.cacheIndex[i]!;
+        const mids: MlxArray[] = [], inputs: MlxArray[] = [];
+        try {
+          for (const group of groups) {
+            const sharedIn = ci === -1 ? group.intermediates[this.previousKvs[i]!]! : null;
+            const { h, shared } = layer.forwardAttn(group.h, group.masks.get(layer.layerType)!,
+              ci === -1 ? null : group.cache[ci]!, sharedIn);
+            mids.push(h); group.intermediates[i] = shared;
+            if (group.perLayer) {
+              const [B, L] = group.ids.shape;
+              using part = group.perLayer.slice([0, 0, i, 0], [B!, L!, i + 1, this.perLayerWidth]);
+              inputs.push(ops.reshape(part, [B!, L!, this.perLayerWidth]));
+            }
+          }
+          // Per-layer inputs use the same token order as the hidden states.
+          let packedInput: MlxArray | null = null;
+          if (packedMlp && inputs.length) {
+            const flat = inputs.map(input => ops.reshape(input, [1, -1, this.perLayerWidth]));
+            try { packedInput = ops.concatAxis(flat, 1); }
+            finally { for (const input of flat) input.dispose(); }
+          }
+          using pli = packedInput;
+          const outputs: MlxArray[] = [];
+          try {
+            if (packedMlp) outputs.push(...mapPackedTokens(mids, packed => layer.forwardMlp(packed, pli)));
+            else for (const [row, mid] of mids.entries()) outputs.push(layer.forwardMlp(mid, inputs[row] ?? null));
+          } catch (error) { for (const output of outputs) output.dispose(); throw error; }
+          for (const [index, group] of groups.entries()) {
+            group.h.dispose(); group.h = outputs[index]!;
+          }
+          for (const [index, group] of groups.entries()) work[index]!.captureLayer?.(i, group.h);
+        } finally {
+          for (const mid of mids) mid.dispose();
+          for (const input of inputs) input.dispose();
+        }
+      }
+      for (const [index, group] of groups.entries()) {
+        const hidden = this.finalNorm.forward(group.h); results.push(hidden);
+        work[index]!.captureLayer?.(this.layers.length, hidden);
+      }
+      return results;
+    } catch (error) { for (const result of results) result.dispose(); throw error; }
+    finally {
+      for (const group of groups) {
+        group.h.dispose(); group.perLayer?.dispose();
+        for (let i = 0; i < this.numDonors; i++) {
+          const shared = group.intermediates[i];
+          if (shared) Gemma4Model.disposeSharedKv(shared);
+        }
+        for (const mask of group.masks.values()) mask.arr?.dispose();
+      }
+    }
+  }
+
+
   // --- Segmented-backward support (docs/design/orpo-training.md
   // §4 Phase B). These are ADDITIVE — forwardLayers is untouched. The segmented
   // driver builds masks + per-layer inputs ONCE, then drives runLayerRange per

package/src/model/kv-attention-view.ts

@@ -2,6 +2,7 @@ import type { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
 import type { Cache, KvAttentionView, KvDonorAttention } from "./gemma4-base";
 import { quantizedSdpa, disposeTriple } from "./gemma4-base";
+import { quantizedAppendAttention } from "./quantized-append-attention";
 
 export function captureKvDonorAttention(cache: Cache): KvDonorAttention {
   if (cache.captureDonorAttention) return cache.captureDonorAttention();
@@ -31,7 +32,9 @@ export function captureKvAttention(cache: Cache, k: MlxArray, v: MlxArray): KvAt
       const { groupSize, bits } = quantized;
       const [keys, values] = quantized.updateAndFetchQuantized(k, v);
       return {
-        attend: (q, scale, mask) => quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
+        attend: (q, scale, mask, independentPositions) => independentPositions && q.shape[2]! > 1
+          ? quantizedAppendAttention(q, keys, values, scale, groupSize, bits)
+          : quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
         dispose() { disposeTriple(keys); disposeTriple(values); },
       };
     }

package/src/model/profile.ts

@@ -7,6 +7,7 @@ import {
   isQwen35Config,
   isQwen3Config,
   isQwen3MoeConfig,
+  isWhisperConfig,
 } from "./support";
 import { GENERATED } from "./generated";
 import { GENERIC_MODEL_TYPES, genericArgsFor, remapModelType } from "./universal/archs";
@@ -16,6 +17,7 @@ export const ENGINE_CAPABILITIES = Object.freeze([
   "colibri-container",
   "diffusion",
   "diffusion-gemma-graph",
+  "encoder-decoder",
   "gemma4-graph",
   "generated-graph",
   "glm5.2-graph",
@@ -30,6 +32,7 @@ export const ENGINE_CAPABILITIES = Object.freeze([
   "streamed-experts",
   "universal-dense-graph",
   "vision-sidecar",
+  "whisper-graph",
 ] as const);
 
 export type EngineCapability = typeof ENGINE_CAPABILITIES[number];
@@ -37,6 +40,7 @@ export type FidelityTier = "l1" | "l2" | "l3";
 
 export type FidelityTarget =
   | Readonly<{ tier: "l1"; oracle: "mlx-lm"; claim: "bit-exact" }>
+  | Readonly<{ tier: "l1"; oracle: "mlx-whisper"; claim: "bit-exact" }>
   | Readonly<{ tier: "l2"; oracle: "mlx-optiq"; claim: "bit-exact" }>
   | Readonly<{ tier: "l3"; oracle: null; claim: "measured" }>;
 
@@ -49,8 +53,9 @@ export type ModelGraph =
   | "qwen3-moe"
   | "diffusion-gemma"
   | "glm5.2"
-  | "universal-dense";
-export type GenerationLoop = "autoregressive" | "diffusion";
+  | "universal-dense"
+  | "whisper";
+export type GenerationLoop = "autoregressive" | "diffusion" | "encoder-decoder";
 export type ModelSpecialization = "artifact" | "dedicated" | "generated" | "generic";
 
 export interface ModelExecutionComposition {
@@ -105,6 +110,11 @@ export interface ResolveModelProfileOptions {
 const L1: FidelityTarget = Object.freeze({
   tier: "l1", oracle: "mlx-lm", claim: "bit-exact",
 });
+/** L1 for speech: the capability's own reference implementation (mlx-whisper)
+ * plays the mlx-lm role — no mlx-lm arm exists for Whisper. */
+const L1_WHISPER: FidelityTarget = Object.freeze({
+  tier: "l1", oracle: "mlx-whisper", claim: "bit-exact",
+});
 const L2: FidelityTarget = Object.freeze({
   tier: "l2", oracle: "mlx-optiq", claim: "bit-exact",
 });
@@ -231,6 +241,14 @@ const FAMILY_PROFILES = {
       loader: "colibri", graph: "glm5.2", loop: "autoregressive", specialization: "dedicated",
     },
   }),
+  whisper: freezeProfile({
+    id: "whisper-dedicated",
+    fidelity: L1_WHISPER,
+    requiredCapabilities: ["safetensors", "encoder-decoder", "whisper-graph"],
+    execution: {
+      loader: "safetensors", graph: "whisper", loop: "encoder-decoder", specialization: "dedicated",
+    },
+  }),
   universal: freezeProfile({
     id: "universal-dense",
     fidelity: L1,
@@ -263,6 +281,7 @@ function familyProfile(config: ModelConfig, fingerprint: string): ModelProfile {
   if (isQwen35Config(config)) return FAMILY_PROFILES.qwen35;
   if (isQwen3MoeConfig(config)) return FAMILY_PROFILES.qwen3Moe;
   if (isQwen3Config(config)) return FAMILY_PROFILES.qwen3;
+  if (isWhisperConfig(config)) return FAMILY_PROFILES.whisper;
   if (config.modelType.startsWith("gemma4"))
     return GENERATED.has(fingerprint) ? FAMILY_PROFILES.gemma4Generated : FAMILY_PROFILES.gemma4;
   if (genericArgsFor(config)) return FAMILY_PROFILES.universal;
@@ -271,7 +290,7 @@ function familyProfile(config: ModelConfig, fingerprint: string): ModelProfile {
   throw new Error(
     `unsupported model_type "${config.modelType}"` +
     (arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
-    ` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, MiniCPM5;` +
+    ` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, whisper, MiniCPM5;` +
     ` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`,
   );
 }
@@ -301,6 +320,7 @@ const GRAPH_METADATA: Readonly<Record<ModelGraph, GraphMetadata>> = Object.freez
     accepts: isGlm52Config,
     capabilities: ["glm5.2-graph", "streamed-experts"],
   },
+  "whisper": { accepts: isWhisperConfig, capabilities: ["whisper-graph"] },
   "universal-dense": {
     accepts: (config) => {
       try { return genericArgsFor(config) !== null; } catch { return false; }
@@ -318,6 +338,7 @@ function executionCapabilities(profile: ModelProfile): EngineCapability[] {
     ? "colibri-container"
     : "safetensors"];
   if (profile.execution.loop === "diffusion") required.push("diffusion");
+  else if (profile.execution.loop === "encoder-decoder") required.push("encoder-decoder");
   else required.push("autoregressive");
   required.push(...GRAPH_METADATA[profile.execution.graph].capabilities);
   if (profile.execution.specialization === "generated") required.push("generated-graph");
@@ -337,6 +358,8 @@ function validateProfile(profile: ModelProfile): void {
   const fidelityOk =
     (profile.fidelity.tier === "l1" && profile.fidelity.oracle === "mlx-lm" &&
       profile.fidelity.claim === "bit-exact") ||
+    (profile.fidelity.tier === "l1" && profile.fidelity.oracle === "mlx-whisper" &&
+      profile.fidelity.claim === "bit-exact" && profile.execution.graph === "whisper") ||
     (profile.fidelity.tier === "l2" && profile.fidelity.oracle === "mlx-optiq" &&
       profile.fidelity.claim === "bit-exact") ||
     (profile.fidelity.tier === "l3" && profile.fidelity.oracle === null &&
@@ -345,7 +368,9 @@ function validateProfile(profile: ModelProfile): void {
   const expectedLoader = profile.execution.graph === "glm5.2" ? "colibri" : "safetensors";
   const expectedLoop = profile.execution.graph === "diffusion-gemma"
     ? "diffusion"
-    : "autoregressive";
+    : profile.execution.graph === "whisper"
+      ? "encoder-decoder"
+      : "autoregressive";
   if (profile.execution.loader !== expectedLoader || profile.execution.loop !== expectedLoop)
     throw new Error(
       `model profile ${profile.id} has an invalid execution composition for ` +