mlx-bun 0.4.0 → 0.5.0

package/src/model/quantized-append-attention.ts

@@ -0,0 +1,31 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import { quantizedSdpa } from "./gemma4-base";
+
+/** Committed positions retain single-token attention arithmetic. Projection
+ * work can still share the span; each query sees its own causal KV prefix. */
+export function quantizedAppendAttention(
+  q: MlxArray, keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
+  scale: number, groupSize: number, bits: number,
+): MlxArray {
+  const [batch, heads, length, dim] = q.shape as [number, number, number, number];
+  const prefix = keys.packed.shape[2]! - length;
+  const outputs: MlxArray[] = [];
+  const trim = (tensor: ops.QuantizedTensor, end: number) => {
+    const slice = (a: MlxArray) => a.slice([0, 0, 0, 0],
+      [a.shape[0]!, a.shape[1]!, end, a.shape[3]!]);
+    return { packed: slice(tensor.packed), scales: slice(tensor.scales), biases: slice(tensor.biases) };
+  };
+  try {
+    for (let row = 0; row < length; row++) {
+      using queryView = q.slice([0, 0, row, 0], [batch, heads, row + 1, dim]);
+      using query = ops.contiguous(queryView);
+      const k = trim(keys, prefix + row + 1), v = trim(values, prefix + row + 1);
+      try { outputs.push(quantizedSdpa(query, k, v, scale, { mode: "", arr: null }, groupSize, bits)); }
+      finally {
+        for (const t of [k, v]) { t.packed.dispose(); t.scales.dispose(); t.biases.dispose(); }
+      }
+    }
+    return ops.concatAxis(outputs, 2);
+  } finally { for (const output of outputs) output.dispose(); }
+}

package/src/model/qwen3-mrope.ts

@@ -30,20 +30,13 @@ export interface MropeRequestState {
 }
 
 export interface MropeForwardState {
-  /** Effective positions [3, 1, L] int32 for the CURRENT forward window. */
+  /** Effective positions [3, B, L] int32 for the current forward window. */
   posIds: MlxArray;
   /** Borrowed model-owned f32 inv_freq (mropeInvFreq) — not disposed here. */
   invFreq: MlxArray;
   rotaryDims: number;
 }
 
-/** Active per-forward mRoPE cos/sin, read by Qwen3Attention.forward. Null on
- *  the text-only path (serial lane only — vision requests never batch). */
-export let activeMrope: MropeForwardState | null = null;
-export function setActiveMrope(s: MropeForwardState | null): void {
-  activeMrope = s;
-}
-
 const MROPE_SECTION = [11, 11, 10] as const;
 
 /** Interleaved frequency→stream selector (rope_utils
@@ -85,6 +78,11 @@ export function buildMropePositions(
   state: MropeRequestState, offset: number, L: number,
   invFreq: MlxArray, rotaryDims: number,
 ): MropeForwardState {
+  return { posIds: mropePositionIds(state, offset, L), invFreq, rotaryDims };
+}
+
+/** Request-owned prompt positions; the returned native array belongs to the caller. */
+export function mropePositionIds(state: MropeRequestState, offset: number, L: number): MlxArray {
   const P = state.positions;
   const promptLen = P[0].length;
   const joint = new Int32Array(3 * L);
@@ -94,7 +92,7 @@ export function buildMropePositions(
       joint[axis * L + l] = idx < promptLen ? P[axis]![idx]! : idx + state.delta;
     }
   }
-  return { posIds: MlxArray.fromInt32(joint, [3, 1, L]), invFreq, rotaryDims };
+  return MlxArray.fromInt32(joint, [3, 1, L]);
 }
 
 // The reference language-side apply is a CUSTOM METAL KERNEL (rope_utils

package/src/model/qwen3_5.ts

@@ -10,8 +10,11 @@
 // shared maybeQuantizeKv path).
 
 import type { ModelConfig } from "../config";
+import { quantizedAppendAttention } from "./quantized-append-attention";
 import { gatedDeltaState } from "./qwen3-delta-state";
 import type { Weights } from "../weights";
+import { runtimeFlag } from "../runtime-config";
+import { mapTokenGroups, type TokenGroup } from "./token-groups";
 import { MlxArray } from "../mlx/array";
 import { Dtype, deviceArchitecture } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
@@ -36,8 +39,8 @@ import { materializeCopy } from "../mlx/materialize";
 import { gatedDeltaUpdate, SSMCache } from "./qwen3-delta";
 import type { QwenConvolution } from "./qwen-conv";
 import {
-  activeMrope, applyInterleavedRope, buildMropePositions, mropeInvFreq,
-  setActiveMrope, type MropeRequestState,
+  applyInterleavedRope, buildMropePositions, mropeInvFreq,
+  type MropeRequestState, type MropeForwardState,
 } from "./qwen3-mrope";
 
 const PREFIX = "language_model";
@@ -406,7 +409,8 @@ export class Qwen3Attention {
     this.kNorm = new RMSNorm(weights.tensor(`${prefix}.k_norm.weight`), t.rmsNormEps);
   }
 
-  forward(x: MlxArray, mask: Mask, cache: Cache, independentRows = false): MlxArray {
+  forward(x: MlxArray, mask: Mask, cache: Cache, independentRows = false,
+    mrope: MropeForwardState | null = null): MlxArray {
     const [B, L] = x.shape as [number, number, number];
 
     // q_proj emits 2× head_dim per head → split into queries + gate.
@@ -433,11 +437,10 @@ export class Qwen3Attention {
     // position as ropeOffsetArr (rows have different prompt lengths); the
     // dynamic-offset kernel is the same fast::rope, bit-exact vs the static
     // form (tests/unit/compile.test.ts). Serial lane: scalar offset, unchanged.
-    // Vision requests (serial lane only) install activeMrope for the current
-    // forward: 3D interleaved positions via the manual apply — the reference
+    // Media forwards supply their own 3D interleaved positions. The reference
     // never uses the fused fast-rope kernel when position_ids are supplied,
     // so this IS the oracle's own arithmetic. Text-only stays on ops.rope.
-    const mr = activeMrope;
+    const mr = mrope;
     const offArr = (cache as { ropeOffsetArr?: MlxArray }).ropeOffsetArr;
     if (mr) {
       q = disposing(q, applyInterleavedRope(q, mr));
@@ -456,12 +459,14 @@ export class Qwen3Attention {
     if (cache.attentionState) {
       const view = cache.attentionState.appendAndFetch(k, v);
       k.dispose(); v.dispose();
-      try { attn = view.attend(q, this.scale, mask); } finally { view.dispose(); }
+      try { attn = view.attend(q, this.scale, mask, independentRows); } finally { view.dispose(); }
     } else if (quantized) {
       const [keys, values] = quantized.updateAndFetchQuantized(k, v);
       k.dispose();
       v.dispose();
-      attn = quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
+      attn = independentRows && L > 1
+        ? quantizedAppendAttention(q, keys, values, this.scale, quantized.groupSize, quantized.bits)
+        : quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
       disposeTriple(keys);
       disposeTriple(values);
     } else {
@@ -559,23 +564,27 @@ export class Qwen3Layer {
     this.postAttnNorm = new RMSNorm(weights.tensor(`${prefix}.post_attention_layernorm.weight`), config.text.rmsNormEps);
   }
 
-  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false, ssmMask?: MlxArray | null): MlxArray {
-    const xn = this.inputNorm.forward(x);
-    const r = this.isLinear
+  forward(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false, ssmMask?: MlxArray | null,
+    mrope: MropeForwardState | null = null): MlxArray {
+    using h = this.forwardAttn(x, faMask, cache, independentRows, ssmMask, mrope);
+    return this.forwardMlp(h, independentRows);
+  }
+
+  forwardAttn(x: MlxArray, faMask: Mask, cache: Cache, independentRows = false, ssmMask?: MlxArray | null,
+    mrope: MropeForwardState | null = null): MlxArray {
+    using xn = this.inputNorm.forward(x);
+    using r = this.isLinear
       ? this.linearAttn!.forward(xn, cache as SSMCache, independentRows, ssmMask)
-      : this.selfAttn!.forward(xn, faMask, cache, independentRows);
-    xn.dispose();
-    const h = ops.add(x, r);
-    r.dispose();
-    const hn = this.postAttnNorm.forward(h);
+      : this.selfAttn!.forward(xn, faMask, cache, independentRows, mrope);
+    return ops.add(x, r);
+  }
+
+  forwardMlp(h: MlxArray, independentRows = false): MlxArray {
+    using hn = this.postAttnNorm.forward(h);
     // RMSNorm allocates an aligned row-contiguous output. Pass that layout
     // proof so Trellis can match native small-prefill matmul without an eval.
-    const m = this.mlp.forward(hn, true, independentRows);
-    hn.dispose();
-    const out = ops.add(h, m);
-    h.dispose();
-    m.dispose();
-    return out;
+    using m = this.mlp.forward(hn, true, independentRows);
+    return ops.add(h, m);
   }
 }
 
@@ -669,10 +678,16 @@ export class Qwen35Model {
         return null;
     }
     return {
-      maxChunkSize: (state: readonly Cache[]) => qwenAppendChunkSize(state[0]!.offset),
+      affineKvBits: [4, 8],
+      turboQuantFormats: [{ kBits: 8, vBits: 3 }],
+      maxChunkSize: (state: readonly Cache[], rows = 1) => rows === 1 ? qwenAppendChunkSize(state[0]!.offset) : 1,
       forwardHidden: (ids: MlxArray, cache: Cache[]): MlxArray => {
+        // Wider cohorts retain their ordinary one-position numerical graph.
+        // The specialized multi-position append remains qualified at B=1.
+        if (ids.shape.length === 2 && ids.shape[0]! > 1 && ids.shape[1] === 1)
+          return this.forwardHidden(ids, cache);
         if (ids.shape.length !== 2 || ids.shape[0] !== 1 || ids.shape[1]! > 4)
-          throw new Error("Qwen committed-token append supports one row and at most four positions");
+          throw new Error("Qwen committed-token append supports one position per shared row or up to four positions at B=1");
         const hidden = this.embed.encode(ids);
         return this.forwardLayers(hidden, cache, true);
       },
@@ -684,10 +699,61 @@ export class Qwen35Model {
     return this.forwardLayers(h, cache);
   }
 
+  /** Borrowed request positions, including different positions for each row. */
+  forwardHiddenAtPositions(ids: MlxArray, cache: Cache[], positions: MlxArray): MlxArray {
+    return this.forwardLayers(this.embed.encode(ids), cache, false, positions);
+  }
+
+  forwardEmbeddingsAtPositions(embeds: MlxArray, cache: Cache[], positions: MlxArray): MlxArray {
+    return this.forwardLayers(ops.contiguous(embeds), cache, false, positions);
+  }
+
+  /** Attention and DeltaNet retain disjoint row state. Only tokenwise MLP
+   * work is eligible for packing; verification can preserve its geometry. */
+  forwardHiddenMixed(work: readonly TokenGroup[]): MlxArray[] {
+    if (work.length === 1 && !work[0]!.captureLayer) return [this.forwardHidden(work[0]!.ids, work[0]!.cache)];
+    const groups: Array<{ h: MlxArray; mask?: Mask; ssmMask?: MlxArray | null }> = [];
+    const results: MlxArray[] = [];
+    const pack = runtimeFlag("MLX_BUN_MIXED_PACKED_MLP", true);
+    const bounded = work.some(group => group.ids.shape[1]! > TRELLIS_MATVEC_MAX_M);
+    try {
+      for (const { ids, cache } of work) {
+        const group: typeof groups[number] = { h: this.embed.encode(ids) }; groups.push(group);
+        group.mask = cache[this.faIdx]!.makeMask(ids.shape[1]!, null);
+        group.ssmMask = (cache[0] as SSMCache).prefillPadding?.makeMask(ids.shape[1]!);
+      }
+      for (const [i, layer] of this.layers.entries()) {
+        const mids: MlxArray[] = [];
+        try {
+          for (const [row, group] of groups.entries())
+            mids.push(layer.forwardAttn(group.h, group.mask!, work[row]!.cache[i]!, false, group.ssmMask));
+          const outputs = mapTokenGroups(work, mids, hidden => layer.forwardMlp(hidden), pack);
+          for (const [row, group] of groups.entries()) { group.h.dispose(); group.h = outputs[row]!; }
+          // Materialize every layer's recurrent tail as well as the residual.
+          // It is not a dependency of h and otherwise pins prefill buffers.
+          if (bounded) {
+            const caches = work.map(group => group.cache[i]!);
+            const state = caches.map(cache => cache.state());
+            try { ops.evalAll([...groups.map(group => group.h), ...state.flat()]); }
+            finally { for (const [row, cache] of caches.entries())
+              if (cache.stateNeedsDispose) for (const value of state[row]!) value.dispose(); }
+          }
+          for (const [row, group] of groups.entries()) work[row]!.captureLayer?.(i, group.h);
+        } finally { for (const mid of mids) mid.dispose(); }
+      }
+      for (const [row, group] of groups.entries()) {
+        const hidden = this.finalNorm.forward(group.h); results.push(hidden);
+        work[row]!.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.mask?.arr?.dispose(); group.ssmMask?.dispose(); } }
+  }
+
   /** Active vision mRoPE request state (serial lane; set by the generation
    *  gateway around a vision request's run, null for text-only — which keeps
-   *  the bit-exact fast-rope path). While set, forwardLayers installs the
-   *  per-forward interleaved cos/sin consumed by every full-attn layer. */
+   *  the bit-exact fast-rope path). The explicit serial executor scopes this value.
+   *  Shared media forwards pass request-owned positions directly. */
   mrope: MropeRequestState | null = null;
   #mropeInvFreq: MlxArray | null = null;
 
@@ -722,7 +788,8 @@ export class Qwen35Model {
     tap.captured.set(i, copy);
   }
 
-  protected forwardLayers(h0: MlxArray, cache: Cache[], independentRows = false): MlxArray {
+  protected forwardLayers(h0: MlxArray, cache: Cache[], independentRows = false,
+    positions?: MlxArray): MlxArray {
     const L = h0.shape[1]!;
     // One full-attention mask shared by all full layers (same offset); linear
     // layers see no ssm mask at B=1.
@@ -731,15 +798,13 @@ export class Qwen35Model {
     // Vision requests: one interleaved-mRoPE cos/sin table per forward,
     // shared by all 12 full-attention layers (positions are layer-invariant).
     let mropeFwd: ReturnType<typeof buildMropePositions> | null = null;
-    if (this.mrope) {
+    if (positions || this.mrope) {
       const t = this.config.text;
       const ropeDims = Math.trunc(t.headDim * t.partialRotaryFactor);
       const base = t.ropeParameters.full_attention?.ropeTheta ?? 10000;
       this.#mropeInvFreq ??= mropeInvFreq(ropeDims, base);
-      mropeFwd = buildMropePositions(
-        this.mrope, cache[this.faIdx]!.offset, L, this.#mropeInvFreq, ropeDims,
-      );
-      setActiveMrope(mropeFwd);
+      mropeFwd = positions ? { posIds: positions, invFreq: this.#mropeInvFreq, rotaryDims: ropeDims }
+        : buildMropePositions(this.mrope!, cache[this.faIdx]!.offset, L, this.#mropeInvFreq, ropeDims);
     }
     let h: MlxArray | null = h0;
     const prof = (globalThis as Record<string, unknown>).__deltaProf as
@@ -748,7 +813,7 @@ export class Qwen35Model {
     try {
       for (let i = 0; i < this.layers.length; i++) {
         const tl = prof ? performance.now() : 0;
-        const next = this.layers[i]!.forward(h, faMask, cache[i]!, independentRows, ssmMask);
+        const next = this.layers[i]!.forward(h, faMask, cache[i]!, independentRows, ssmMask, mropeFwd);
         h.dispose();
         h = next;
         if (L > TRELLIS_MATVEC_MAX_M) {
@@ -772,10 +837,7 @@ export class Qwen35Model {
       h = null; // consumed — the finally must not double-free
       return out;
     } finally {
-      if (mropeFwd) {
-        setActiveMrope(null);
-        mropeFwd.posIds.dispose();
-      }
+      if (mropeFwd && !positions) mropeFwd.posIds.dispose();
       // A mid-loop layer throw must not strand the mask array or the
       // in-flight [1,L,H] residual (2026-08-18 review).
       faMask.arr?.dispose();

package/src/model/support.ts

@@ -61,6 +61,13 @@ export function isDiffusionGemmaConfig(config: ModelConfig): boolean {
   return config.modelType === "diffusion_gemma";
 }
 
+/** Whisper (model_type `whisper`, mlx-whisper ModelDimensions config):
+ *  encoder-decoder speech recognition — audio in, text out. Routed through
+ *  the transcription engine (src/audio/whisper-*.ts), never the chat loop. */
+export function isWhisperConfig(config: ModelConfig): boolean {
+  return config.modelType === "whisper" && typeof config.raw.n_audio_ctx === "number";
+}
+
 /** Support tier for a registry record (docs/design/generic-model-support.md):
  *  "targeted" = dedicated/generated forward + L2/L3 paths;
  *  "generic"  = the Tier-0 universal module (L1 monolith only);
@@ -74,6 +81,7 @@ export function supportTier(modelType: string, repoId = ""): "targeted" | "gener
   if (modelType === "qwen3") return "targeted"; // plain Qwen3 (embedding backbone)
   if (modelType === "qwen3_moe") return "targeted"; // Qwen3-MoE (sparse experts)
   if (modelType === "glm_moe_dsa") return "targeted"; // dedicated GLM-5.2 MLA/DSA port
+  if (modelType === "whisper") return "targeted"; // speech-to-text (transcription engine)
   if (modelType === "llama" && repoId.toLowerCase().includes("minicpm5-1b-optiq-4bit"))
     return "targeted";
   // Tier-0 generic fallback: the universal-dense descriptor table.

package/src/model/token-groups.ts

@@ -0,0 +1,66 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { Cache } from "./gemma4-base";
+
+export interface TokenWorkOptions {
+  /** Verification keeps its established GEMV/GEMM geometry. Other groups may
+   * still pack their tokenwise work independently. */
+  readonly preserveTokenGeometry?: boolean;
+  /** Borrows a residual stream; the method copies only the layers it needs.
+   * The layer-count index denotes the post-final-norm output. */
+  readonly captureLayer?: (layer: number, hidden: MlxArray) => void;
+}
+
+/** Independent sequences presented to one model execution. Each item keeps
+ * its existing rectangular row geometry and owns a disjoint cache. */
+export interface TokenGroup extends TokenWorkOptions {
+  readonly ids: MlxArray;
+  readonly cache: Cache[];
+}
+
+/** Apply a tokenwise block while respecting each method's numerical policy. */
+export function mapTokenGroups(
+  groups: readonly TokenGroup[], inputs: readonly MlxArray[],
+  operation: (hidden: MlxArray) => MlxArray, pack = true,
+): MlxArray[] {
+  const results: MlxArray[] = [];
+  const packedRows = groups.flatMap((group, row) => pack && !group.preserveTokenGeometry ? [row] : []);
+  try {
+    const packed = packedRows.length ? mapPackedTokens(packedRows.map(row => inputs[row]!), operation) : [];
+    for (const [index, row] of packedRows.entries()) results[row] = packed[index]!;
+    for (let row = 0; row < groups.length; row++) results[row] ??= operation(inputs[row]!);
+    return results;
+  } catch (error) { for (const result of results) result?.dispose(); throw error; }
+}
+
+export interface MixedTokenModel {
+  /** Borrows inputs/state; returns one owned hidden array per input group. */
+  forwardHiddenMixed(groups: readonly TokenGroup[]): MlxArray[];
+}
+
+/** Pack real tokens for a tokenwise operation, then restore each group's
+ * geometry. Attention and recurrent operations must run outside this port.
+ * No padding, host readback, or evaluation is introduced here. */
+export function mapPackedTokens(
+  inputs: readonly MlxArray[], operation: (packed: MlxArray) => MlxArray,
+): MlxArray[] {
+  if (inputs.length === 1) return [operation(inputs[0]!)];
+  const flat = inputs.map(input => ops.reshape(input, [1, -1, input.shape.at(-1)!]));
+  let packed: MlxArray;
+  try { packed = ops.concatAxis(flat, 1); }
+  finally { for (const input of flat) input.dispose(); }
+  using joined = packed;
+  using output = operation(joined);
+  const results: MlxArray[] = [];
+  let position = 0;
+  try {
+    for (const input of inputs) {
+      const [batch, length] = input.shape;
+      const end = position + batch! * length!;
+      using part = output.slice([0, position, 0], [1, end, output.shape[2]!]);
+      results.push(ops.reshape(part, [batch!, length!, output.shape[2]!]));
+      position = end;
+    }
+    return results;
+  } catch (error) { for (const result of results) result.dispose(); throw error; }
+}

package/src/model/trellis-linear.ts

@@ -41,7 +41,7 @@ import * as ops from "../mlx/ops";
 import { MetalKernel } from "../mlx/metal-kernel";
 import type { Weights } from "../weights";
 import { quantFor, type ModelConfig, type QuantSpec } from "../config";
-import { runtimeFlag, runtimeNumber } from "../runtime-config";
+import { runtimeFlag, runtimeNumber, runtimeValue } from "../runtime-config";
 import { lut1mad, wordsPerBlock } from "../quantize/trellis";
 import { QuantizedLinear } from "./gemma4-base";
 import { REDUCE_SHARED_M_SOURCE, GATEUP_SHARED_M_SOURCE } from "./trellis-shared-m";
@@ -475,7 +475,7 @@ export function expandTrellis(codes: MlxArray, scales: MlxArray, g: TrellisGeome
 export type TrellisMode = "kernel" | "expand";
 
 export function trellisModeFromEnv(): TrellisMode {
-  const v = process.env.MLX_BUN_TRELLIS;
+  const v = runtimeValue("MLX_BUN_TRELLIS");
   if (v === "expand") return "expand";
   return "kernel";
 }

package/src/model/turboquant-codec.ts

@@ -1,7 +1,13 @@
+import { runtimeConfig, type RuntimeConfig } from "../runtime-config";
 import type { MlxArray } from "../mlx/array";
 import * as tq from "../mlx/turboquant-ops";
 import type { KvCodec } from "../backends/mlx/kv-codec";
 
+/** Resolve once when a cache is created; storage copies retain this policy. */
+export function turboQuantFusedDecode(config: RuntimeConfig = runtimeConfig()): boolean {
+  return config.flag("MLX_BUN_TURBOQUANT_FUSED_DECODE", true);
+}
+
 /** One TurboQuant-encoded (K, V) storage tuple — the 5 arrays kv-store.ts
  *  and toQuantized/fromKVCache pass around. Not `ops.QuantizedTensor`
  *  (mlx's affine int4/int8 scheme): this is the rotation + Lloyd-Max