mlx-bun 0.4.0 → 0.5.0

package/src/storage/kv-writer.worker.js

@@ -0,0 +1,264 @@
+// CPU-only persistence worker. It deliberately imports no model, MLX,
+// scheduler or cache manager. The sender owns all buffers through completion.
+import { parentPort } from "node:worker_threads";
+import { toArrayBuffer, ptr } from "bun:ffi";
+import { openSync, writeSync, closeSync, fsyncSync, renameSync, rmSync, mkdirSync, existsSync, readSync, readFileSync, readdirSync } from "node:fs";
+import { createHash } from "node:crypto";
+import { dirname, join } from "node:path";
+
+const MAGIC = "MLXBUNKV2\n";
+const PREFIX_LEN = MAGIC.length + 16;
+const alignUp = n => Math.ceil(n / 16384) * 16384;
+const hash64 = bytes => Bun.hash(bytes).toString(16).padStart(16, "0");
+let scratch = new Uint8Array(0);
+
+function bytesOf(tensor) {
+  const { pointer, shape, strides, itemSize } = tensor;
+  const size = shape.reduce((a, b) => a * b, 1);
+  if (!size) return new Uint8Array(0);
+  // Coalesce the contiguous suffix. Most KV layouts need only a copy per
+  // head to omit capacity padding, never a GPU contiguous() operation.
+  let run = 1, axis = shape.length - 1;
+  while (axis >= 0 && (shape[axis] === 1 || strides[axis] === run)) {
+    run *= shape[axis--];
+  }
+  if (axis < 0) return new Uint8Array(toArrayBuffer(pointer, 0, size * itemSize));
+  let lo = 0, hi = 0;
+  for (let i = 0; i < shape.length; i++) {
+    const end = (shape[i] - 1) * strides[i];
+    lo += Math.min(0, end); hi += Math.max(0, end);
+  }
+  const source = new Uint8Array(toArrayBuffer(pointer + lo * itemSize, 0, (hi - lo + 1) * itemSize));
+  if (scratch.length < size * itemSize) scratch = new Uint8Array(size * itemSize);
+  const bytes = scratch.subarray(0, size * itemSize);
+  const runBytes = run * itemSize;
+  for (let row = 0; row < size / run; row++) {
+    let index = row, offset = -lo;
+    for (let i = axis; i >= 0; i--) {
+      offset += (index % shape[i]) * strides[i];
+      index = Math.floor(index / shape[i]);
+    }
+    bytes.set(source.subarray(offset * itemSize, offset * itemSize + runBytes), row * runBytes);
+  }
+  return bytes;
+}
+
+function writeAll(fd, bytes, position) {
+  let offset = 0;
+  while (offset < bytes.length) {
+    const n = writeSync(fd, bytes, offset, bytes.length - offset, position + offset);
+    if (!n) throw new Error("KV write made no progress");
+    offset += n;
+  }
+}
+
+function writeSnapshot({ path, header, tensors }) {
+  const slots = [...header.caches, ...(header.attachments ?? [])].flatMap(entry => entry.tensors);
+  const headerLen = new TextEncoder().encode(JSON.stringify(header)).length;
+  const dataStart = alignUp(PREFIX_LEN + headerLen);
+  mkdirSync(dirname(path), { recursive: true });
+  const tmp = `${path}.tmp`;
+  const fd = openSync(tmp, "w");
+  try {
+    for (let i = 0; i < slots.length; i++) {
+      const bytes = bytesOf(tensors[i]);
+      if (bytes.length !== slots[i].bytes) throw new Error("KV tensor byte-length drift");
+      slots[i].hash = hash64(bytes);
+      writeAll(fd, bytes, dataStart + slots[i].off);
+    }
+    const json = new TextEncoder().encode(JSON.stringify(header));
+    if (json.length !== headerLen) throw new Error("KV header byte-length drift");
+    const prefix = new Uint8Array(dataStart);
+    prefix.set(new TextEncoder().encode(MAGIC));
+    const view = new DataView(prefix.buffer);
+    view.setUint32(MAGIC.length, headerLen, true);
+    view.setUint32(MAGIC.length + 4, dataStart, true);
+    view.setBigUint64(MAGIC.length + 8, BigInt(Bun.hash(json)), true);
+    prefix.set(json, PREFIX_LEN);
+    writeAll(fd, prefix, 0);
+    fsyncSync(fd);
+  } catch (error) {
+    closeSync(fd);
+    rmSync(tmp, { force: true });
+    throw error;
+  }
+  closeSync(fd);
+  try { renameSync(tmp, path); }
+  catch (error) { rmSync(tmp, { force: true }); throw error; }
+}
+
+// The block boundary resets for each leading row (e.g. attention head).
+// Growing a token axis therefore leaves preceding blocks byte-identical.
+function* segments(tensor, blockBytes, metrics) {
+  const { pointer, shape, strides, itemSize } = tensor;
+  const size = shape.reduce((a, b) => a * b, 1);
+  const rowSize = shape.slice(-2).reduce((a, b) => a * b, 1);
+  if (!size) return;
+  let run = 1, axis = shape.length - 1;
+  while (axis >= 0 && (shape[axis] === 1 || strides[axis] === run)) run *= shape[axis--];
+  const limit = Math.max(1, Math.floor(blockBytes / itemSize));
+  let packing = null;
+  for (let start = 0; start < size;) {
+    const count = Math.min(limit, rowSize - start % rowSize, size - start);
+    let index = start, offset = 0;
+    for (let i = shape.length - 1; i >= 0; i--) {
+      offset += index % shape[i] * strides[i]; index = Math.floor(index / shape[i]);
+    }
+    if (count <= run - start % run) {
+      yield new Uint8Array(toArrayBuffer(pointer + offset * itemSize, 0, count * itemSize));
+    } else {
+      if (!packing) packing = new Uint8Array(limit * itemSize);
+      metrics.scratchPeak = Math.max(metrics.scratchPeak, packing.byteLength);
+      const bytes = packing.subarray(0, count * itemSize);
+      for (let copied = 0; copied < count;) {
+        index = start + copied; offset = 0;
+        for (let i = shape.length - 1; i >= 0; i--) {
+          offset += index % shape[i] * strides[i]; index = Math.floor(index / shape[i]);
+        }
+        const n = Math.min(count - copied, run - (start + copied) % run);
+        bytes.set(new Uint8Array(toArrayBuffer(pointer + offset * itemSize, 0, n * itemSize)), copied * itemSize);
+        copied += n;
+      }
+      metrics.copiedBytes += bytes.length;
+      yield bytes;
+    }
+    start += count;
+  }
+}
+
+function writeHeader(path, header) {
+  const json = new TextEncoder().encode(JSON.stringify(header));
+  const prefix = new Uint8Array(PREFIX_LEN + json.length);
+  prefix.set(new TextEncoder().encode(MAGIC));
+  const view = new DataView(prefix.buffer);
+  view.setUint32(MAGIC.length, json.length, true);
+  view.setUint32(MAGIC.length + 4, prefix.length, true);
+  view.setBigUint64(MAGIC.length + 8, BigInt(Bun.hash(json)), true);
+  prefix.set(json, PREFIX_LEN);
+  const tmp = path + ".tmp", fd = openSync(tmp, "w");
+  try { writeAll(fd, prefix, 0); fsyncSync(fd); }
+  finally { closeSync(fd); }
+  renameSync(tmp, path);
+  return prefix.length;
+}
+
+function writeBlocks({ path, header, tensors, blockBytes = 1024 * 1024, segmented = true }) {
+  const metrics = { writtenBytes: 0, reusedBytes: 0, copiedBytes: 0, scratchPeak: 0 };
+  const root = join(dirname(path), "blocks");
+  mkdirSync(root, { recursive: true });
+  const slots = [...header.caches, ...(header.attachments ?? [])].flatMap(entry => entry.tensors);
+  for (let i = 0; i < slots.length; i++) {
+    const slot = slots[i]; slot.blocks = [];
+    let bytes = 0;
+    let source = tensors[i];
+    let packed;
+    if (!segmented) {
+      packed = bytesOf(source);
+      metrics.copiedBytes += packed.buffer === scratch.buffer ? packed.length : 0;
+      metrics.scratchPeak = Math.max(metrics.scratchPeak, scratch.length);
+      // Contiguous strides, but retain row boundaries for identical blocks.
+      let stride = 1;
+      const strides = Array(source.shape.length);
+      for (let j = strides.length - 1; j >= 0; j--) { strides[j] = stride; stride *= source.shape[j]; }
+      // No second copy; the pointer remains live until this tensor finishes.
+      source = { ...source, pointer: packed.length ? Number(ptr(packed)) : 0, strides };
+    }
+    for (const chunk of segments(source, blockBytes, metrics)) {
+      const hash = createHash("sha256").update(chunk).digest("hex");
+      bytes += chunk.length;
+      slot.blocks.push({ hash, bytes: chunk.length });
+      const destination = join(root, hash);
+      if (existsSync(destination)) { metrics.reusedBytes += chunk.length; continue; }
+      const tmp = destination + ".tmp", fd = openSync(tmp, "w");
+      try { writeAll(fd, chunk, 0); fsyncSync(fd); }
+      finally { closeSync(fd); }
+      renameSync(tmp, destination);
+      metrics.writtenBytes += chunk.length;
+    }
+    if (bytes !== slot.bytes) throw new Error("KV tensor byte-length drift");
+    // The header authenticates block order; each block authenticates bytes.
+    // Avoid hashing the same tensor a second time.
+    slot.hash = "";
+  }
+  header.formatVersion = 5;
+  metrics.writtenBytes += writeHeader(path, header);
+  return metrics;
+}
+
+function readAll(fd, bytes, position) {
+  let offset = 0;
+  while (offset < bytes.length) {
+    const n = readSync(fd, bytes, offset, bytes.length - offset, position + offset);
+    if (!n) throw new Error("truncated KV tensor");
+    offset += n;
+  }
+}
+
+// Destinations are unpublished, independently owned host allocations. No
+// MLX API is called here, and the owner retains them until acknowledgement.
+function readSnapshot({ path, header, pointers, verify }) {
+  const slots = [...header.caches, ...(header.attachments ?? [])].flatMap(entry => entry.tensors);
+  const fd = header.formatVersion === 5 ? null : openSync(path, "r");
+  try {
+    for (let i = 0; i < slots.length; i++) {
+      const slot = slots[i];
+      const bytes = slot.bytes ? new Uint8Array(toArrayBuffer(pointers[i], 0, slot.bytes)) : new Uint8Array(0);
+      if (header.formatVersion === 5) {
+        let offset = 0;
+        for (const block of slot.blocks ?? []) {
+          if (!/^[a-f0-9]{64}$/.test(block.hash) || block.bytes <= 0 || offset + block.bytes > bytes.length)
+            throw new Error("invalid KV block reference");
+          const blockFd = openSync(join(dirname(path), "blocks", block.hash), "r");
+          const target = bytes.subarray(offset, offset + block.bytes);
+          try { readAll(blockFd, target, 0); } finally { closeSync(blockFd); }
+          if (verify && createHash("sha256").update(target).digest("hex") !== block.hash)
+            throw new Error("KV block hash mismatch");
+          offset += block.bytes;
+        }
+        if (offset !== bytes.length) throw new Error("incomplete KV tensor blocks");
+      } else {
+        readAll(fd, bytes, header.dataStart + slot.off);
+        if (verify && hash64(bytes) !== slot.hash) throw new Error("KV tensor hash mismatch");
+      }
+    }
+  } finally { if (fd !== null) closeSync(fd); }
+}
+
+// Serialized with writes: an in-flight manifest cannot lose its blocks.
+// Only unreferenced content is collected; shared prefixes remain on SSD.
+function collectBlocks({ root }) {
+  if (!existsSync(root)) return;
+  for (const ns of readdirSync(root, { withFileTypes: true })) {
+    if (!ns.isDirectory()) continue;
+    const dir = join(root, ns.name), blocks = join(dir, "blocks");
+    if (!existsSync(blocks)) continue;
+    const live = new Set();
+    for (const file of readdirSync(dir)) {
+      if (!file.endsWith(".mlxkv")) continue;
+      const fd = openSync(join(dir, file), "r");
+      try {
+        const prefix = new Uint8Array(PREFIX_LEN); readAll(fd, prefix, 0);
+        const view = new DataView(prefix.buffer), json = new Uint8Array(view.getUint32(MAGIC.length, true));
+        readAll(fd, json, PREFIX_LEN);
+        if (BigInt(Bun.hash(json)) !== view.getBigUint64(MAGIC.length + 8, true)) continue;
+        const header = JSON.parse(new TextDecoder().decode(json));
+        for (const entry of [...header.caches, ...(header.attachments ?? [])])
+          for (const slot of entry.tensors) for (const block of slot.blocks ?? []) live.add(block.hash);
+      } finally { closeSync(fd); }
+    }
+    for (const file of readdirSync(blocks))
+      if (!live.has(file)) rmSync(join(blocks, file), { force: true });
+  }
+}
+
+parentPort.on("message", data => {
+  try {
+    const started = performance.now();
+    const metrics = data.operation === "read" ? readSnapshot(data)
+      : data.operation === "collect" ? collectBlocks(data)
+      : data.layout === "blocks" ? writeBlocks(data) : writeSnapshot(data);
+    parentPort.postMessage({ id: data.id, metrics: { ...metrics, elapsedMs: performance.now() - started } });
+  } catch (error) {
+    parentPort.postMessage({ id: data.id, error: String(error) });
+  }
+});

package/src/storage/retention-policy.ts

@@ -0,0 +1,37 @@
+/** RAM policy sees reuse metadata, never tensors, storage or scheduling. */
+export interface RetentionCandidate {
+  bytes: number;
+  cost: number;
+  lastUsed: number;
+  uses: number;
+  priority: number;
+}
+export interface RetentionPolicy {
+  readonly name: string;
+  accessed(entry: RetentionCandidate): void;
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T;
+  evicted(entry: RetentionCandidate): void;
+}
+export class LruRetention implements RetentionPolicy {
+  readonly name = "lru";
+  accessed(_entry: RetentionCandidate): void {}
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T {
+    return entries.reduce((a, b) => a.lastUsed < b.lastUsed ? a : b);
+  }
+  evicted(_entry: RetentionCandidate): void {}
+}
+
+/** GreedyDual size/frequency: saved prefill tokens are the cost proxy.
+ * Advancing age on eviction lets old popularity decay without wall timers. */
+export class CostSizeRetention implements RetentionPolicy {
+  readonly name = "cost-size";
+  #age = 0;
+  accessed(entry: RetentionCandidate): void {
+    entry.priority = this.#age + entry.cost * entry.uses / Math.max(1, entry.bytes);
+  }
+  victim<T extends RetentionCandidate>(entries: readonly T[]): T {
+    return entries.reduce((a, b) => a.priority < b.priority ||
+      (a.priority === b.priority && a.lastUsed < b.lastUsed) ? a : b);
+  }
+  evicted(entry: RetentionCandidate): void { this.#age = Math.max(this.#age, entry.priority); }
+}

package/src/tiered-prompt-cache.ts

@@ -0,0 +1,45 @@
+import { PromptCache, cacheBytes, type ColdTier } from "./prompt-cache";
+import { cloneKvCaches, SpillQueue } from "./kv-store";
+import { SsdDurabilityCoordinator } from "./ssd-durability";
+import { disposeResources } from "./engine/resources";
+import type { SsdCacheStore } from "./ssd-cache";
+
+/** One cache with RAM and SSD storage. Inference sees the same put/take
+ * interface. Recency and capacity choose RAM residents; a pending write is
+ * never treated as a completed demotion. */
+export class TieredPromptCache extends PromptCache {
+  protected override promoteRestores = true;
+  readonly spillQueue: SpillQueue;
+  readonly durability: SsdDurabilityCoordinator;
+
+  constructor(
+    maxBytes: number,
+    storage: Pick<SsdCacheStore, "storeAsync" | "hasDurablePrefix">,
+    cold: ColdTier,
+    clone = cloneKvCaches,
+    readonly writeBehind = true,
+  ) {
+    super(maxBytes, null, cold, clone);
+    this.spillQueue = new SpillQueue(Infinity, cacheBytes, async item => {
+      return storage.hasDurablePrefix(item.tokens, item.ns) ||
+        await storage.storeAsync(item.tokens, item.caches, item.ns, undefined, item.attachments);
+    }, disposeResources, () => false);
+    this.durability = new SsdDurabilityCoordinator(this, this.spillQueue, clone,
+      (tokens, ns) => storage.hasDurablePrefix(tokens, ns), 0, 5_000, () => {
+        // Writing alone does not remove a RAM entry. Apply capacity and
+        // pressure policy after the completed writer releases its views.
+        this.evictToBudget();
+        this.reclaim();
+      });
+    this.canEvict = entry => {
+      if (storage.hasDurablePrefix(entry.tokens, entry.ns)) return true;
+      this.durability.schedule(entry.tokens, entry.ns, false);
+      return false;
+    };
+  }
+
+  override put(...args: Parameters<PromptCache["put"]>): void {
+    super.put(...args);
+    if (this.writeBehind) this.durability.schedule(args[0], args[2]);
+  }
+}

package/src/vision/embedder.ts

@@ -7,6 +7,7 @@
 //
 // Sidecar weights stay bf16; features come out bf16 like text embeddings.
 
+import { tensorFingerprint } from "../model/fingerprint";
 import { ptr, read } from "bun:ffi";
 import { MlxArray, cpuStream } from "../mlx/array";
 import { C } from "../mlx/ffi";
@@ -18,6 +19,13 @@ const cstr = (s: string) => Buffer.from(s + "\0", "utf8");
 
 export class VisionTower {
   #weights = new Map<string, MlxArray>();
+  #cacheIdentity?: string;
+  /** Immutable encoder weights and computation settings, hashed once on use. */
+  get cacheIdentity(): string {
+    return this.#cacheIdentity ??= new Bun.CryptoHasher("sha256")
+      .update("gemma-vision-v1").update(JSON.stringify([this.embedScale, this.rmsNormEps]))
+      .update(tensorFingerprint(this.#weights)).digest("hex");
+  }
   readonly embedScale: number;
   readonly rmsNormEps: number;
 

package/src/vision/prompt.ts

@@ -17,6 +17,8 @@
 //     multimodal soft tokens (image | audio) — multimodalMask, decoupled
 //     from the bidirectional mask.
 
+import type { EncoderCache } from "../backends/mlx/encoder-cache";
+import type { PromptNativeWork } from "../serve/model-binding";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
@@ -35,6 +37,7 @@ import { fetchMediaBytes, videoMediaFetchPolicy } from "../media-fetch";
  *  embed_scale). The two own different preprocessing, so the tower carries
  *  its own. */
 export interface VisionEncoder<P extends { softTokens: number } = { softTokens: number }> {
+  readonly cacheIdentity?: string;
   preprocess(bytes: Uint8Array): Promise<P>;
   features(pre: P): MlxArray;
   dispose?(): void;
@@ -42,6 +45,7 @@ export interface VisionEncoder<P extends { softTokens: number } = { softTokens:
 
 export interface VisionPrompt {
   ids: number[];
+  prefixIdentity?: string;
   /** Unscaled merged embeddings [1, L, hidden] — caller disposes. */
   embeddings: MlxArray;
   /** bool [L] image-token mask — caller disposes. */
@@ -50,6 +54,8 @@ export interface VisionPrompt {
 
 export interface MultimodalPrompt {
   ids: number[];
+  /** Includes every media item and the exact token prefix through the last one. */
+  prefixIdentity?: string;
   /** Unscaled merged embeddings [1, L, hidden] — caller disposes. */
   embeddings: MlxArray;
   /** bool [L] image-token mask for the bidirectional attention overlay.
@@ -76,8 +82,8 @@ export interface AudioTokenIds {
 /** The towers a multimodal prompt may need; pass only what the request uses
  *  (vision-only and audio-only fall out as special cases). */
 export interface MultimodalTowers<P extends { softTokens: number }> {
-  vision?: { tower: VisionEncoder<P>; tokenIds: VisionTokenIds };
-  audio?: { tower: AudioTower; tokenIds: AudioTokenIds };
+  vision?: { tower: VisionEncoder<P>; tokenIds: VisionTokenIds; cache?: EncoderCache };
+  audio?: { tower: AudioTower; tokenIds: AudioTokenIds; cache?: EncoderCache };
 }
 
 /** Extract image bytes from OpenAI-style content parts, rewriting the
@@ -223,6 +229,7 @@ export async function buildMultimodalPrompt<P extends { softTokens: number }>(
   images: Uint8Array[],
   audio: Uint8Array[],
   tools: ToolDefinition[] | null = null,
+  nativeWork: PromptNativeWork = (work) => work(),
 ): Promise<MultimodalPrompt> {
   if (images.length > 0 && !towers.vision)
     throw new Error("images present but no vision tower configured");
@@ -280,70 +287,97 @@ export async function buildMultimodalPrompt<P extends { softTokens: number }>(
   if (audIdx !== audio.length)
     throw new Error(`${audio.length} audio clips but ${audIdx} <|audio|> markers`);
 
-  // embed text tokens, then overwrite soft-token rows with tower features
-  const idsArr = ops.fromInt32(spliced, [1, spliced.length]);
-  let embeds = model.embed.encode(idsArr);
-  idsArr.dispose();
-  const hidden = embeds.shape[2]!;
-  for (const run of runs) {
-    let cast: MlxArray;
-    if (run.kind === "image") {
-      // vision towers return pre-divided by embed_scale (their oracle's
-      // frontend convention) — cast to the embedding dtype and splice.
-      const feats = towers.vision!.tower.features(pre[run.index]!); // [1, soft, hidden]
-      cast = feats.astype(embeds.dtype);
-      feats.dispose();
-    } else {
-      // audio mirrors gen-e4b-audio-golden.py EXACTLY:
-      //   features.astype(embeds.dtype) / embed_scale
-      // — raw f32 embed_audio output, cast to bf16 FIRST, then divided by a
-      // weak (dtype-following, i.e. bf16) embed_scale scalar. The order is
-      // load-bearing for the bit-exact greedy gate (see AudioTower.features).
-      const clip = clips[run.index]!;
-      const raw = towers.audio!.tower.features(clip.mel, clip.frames, false); // [1, n, hidden] f32
-      if (raw.shape[1] !== clip.softTokens) {
-        const n = raw.shape[1];
-        raw.dispose();
-        embeds.dispose();
-        throw new Error(`audio tower produced ${n} frames but splice expects ${clip.softTokens}`);
-      }
-      const bf = raw.astype(embeds.dtype);
-      raw.dispose();
-      const scale = ops.scalarLike(towers.audio!.tower.embedScale, bf);
-      cast = ops.div(bf, scale);
-      bf.dispose();
-      scale.dispose();
+  // Exact encoded objects contain no request positions or template state.
+  // CPU preprocessing and SSD reads run before acquiring native execution.
+  const keys = runs.map(run => towers[run.kind === "image" ? "vision" : "audio"]?.cache
+    ? new Bun.CryptoHasher("sha256").update(`gemma-${run.kind}-input-v1`)
+      .update(run.kind === "image" ? images[run.index]! : audio[run.index]!).digest("hex")
+    : "");
+  const lastMedia = runs.at(-1);
+  // A matching namespace guarantees any reused prefix includes all media.
+  // The audio-presence bit also identifies Gemma's causal/bidirectional policy.
+  const prefixIdentity = lastMedia && keys.every(Boolean)
+    ? new Bun.CryptoHasher("sha256").update("gemma-prepared-prefix-v1")
+      .update(JSON.stringify({ audio: audio.length > 0,
+        tokens: spliced.slice(0, lastMedia.start + lastMedia.length),
+        media: runs.map((run, i) => [run.kind, keys[i],
+          (run.kind === "image" ? towers.vision : towers.audio)!.cache!.identity]),
+      })).digest("hex") : undefined;
+  const reused: (MlxArray | null)[] = [];
+  try {
+    for (const [index, run] of runs.entries()) {
+      const cache = run.kind === "image" ? towers.vision?.cache : towers.audio?.cache;
+      reused.push(cache ? await cache.take(keys[index]!) : null);
     }
-    const { start, length } = run;
-    const updated = ops.sliceUpdate(embeds, cast, [0, start, 0], [1, start + length, hidden]);
-    cast.dispose();
-    embeds.dispose();
-    embeds = updated;
-  }
+    return await nativeWork(async () => {
+      // embed text tokens, then overwrite soft-token rows with tower features
+      using idsArr = ops.fromInt32(spliced, [1, spliced.length]);
+      let embeds = model.embed.encode(idsArr);
+      try {
+        const hidden = embeds.shape[2]!;
+        for (const [runIndex, run] of runs.entries()) {
+          let cast: MlxArray;
+          if (run.kind === "image") {
+            // vision towers return pre-divided by embed_scale (their oracle's
+            // frontend convention) — cast to the embedding dtype and splice.
+            const cached = reused[runIndex];
+            using feats = cached ?? towers.vision!.tower.features(pre[run.index]!); // [1, soft, hidden]
+            reused[runIndex] = null;
+            if (!cached) towers.vision!.cache?.put(keys[runIndex]!, feats);
+            cast = feats.astype(embeds.dtype);
+          } else {
+            // audio mirrors gen-e4b-audio-golden.py EXACTLY:
+            //   features.astype(embeds.dtype) / embed_scale
+            // — raw f32 embed_audio output, cast to bf16 FIRST, then divided by a
+            // weak (dtype-following, i.e. bf16) embed_scale scalar. The order is
+            // load-bearing for the bit-exact greedy gate (see AudioTower.features).
+            const clip = clips[run.index]!;
+            const cached = reused[runIndex];
+            using raw = cached ?? towers.audio!.tower.features(clip.mel, clip.frames, false); // [1, n, hidden] f32
+            reused[runIndex] = null;
+            if (raw.shape[1] !== clip.softTokens) {
+              const n = raw.shape[1];
+              throw new Error(`audio tower produced ${n} frames but splice expects ${clip.softTokens}`);
+            }
+            if (!cached) towers.audio!.cache?.put(keys[runIndex]!, raw);
+            using bf = raw.astype(embeds.dtype);
+            using scale = ops.scalarLike(towers.audio!.tower.embedScale, bf);
+            cast = ops.div(bf, scale);
+          }
+          const { start, length } = run;
+          try {
+            const updated = ops.sliceUpdate(embeds, cast, [0, start, 0], [1, start + length, hidden]);
+            embeds.dispose();
+            embeds = updated;
+          } finally { cast.dispose(); }
+        }
 
-  // masks over the spliced sequence: union multimodal soft tokens (per-layer
-  // id zeroing) always; image-only bidirectional overlay only when NO audio
-  // is present (§3.3 Q1).
-  const mmInts = new Int32Array(spliced.length);
-  for (let i = 0; i < spliced.length; i++) {
-    const id = spliced[i]!;
-    mmInts[i] = (vIds && id === vIds.imageTokenId) || (aIds && id === aIds.audioTokenId) ? 1 : 0;
-  }
-  const mmI32 = MlxArray.fromInt32(mmInts, [spliced.length]);
-  const multimodalMask = mmI32.astype(Dtype.bool);
-  mmI32.dispose();
+        // masks over the spliced sequence: union multimodal soft tokens (per-layer
+        // id zeroing) always; image-only bidirectional overlay only when NO audio
+        // is present (§3.3 Q1).
+        const mmInts = new Int32Array(spliced.length);
+        for (let i = 0; i < spliced.length; i++) {
+          const id = spliced[i]!;
+          mmInts[i] = (vIds && id === vIds.imageTokenId) || (aIds && id === aIds.audioTokenId) ? 1 : 0;
+        }
+        const mmI32 = MlxArray.fromInt32(mmInts, [spliced.length]);
+        const multimodalMask = mmI32.astype(Dtype.bool);
+        mmI32.dispose();
 
-  let bidirMask: MlxArray | null = null;
-  if (images.length > 0 && audio.length === 0) {
-    const imgInts = new Int32Array(spliced.length);
-    for (let i = 0; i < spliced.length; i++)
-      imgInts[i] = spliced[i] === vIds!.imageTokenId ? 1 : 0;
-    const imgI32 = MlxArray.fromInt32(imgInts, [spliced.length]);
-    bidirMask = imgI32.astype(Dtype.bool);
-    imgI32.dispose();
-  }
+        let bidirMask: MlxArray | null = null;
+        if (images.length > 0 && audio.length === 0) {
+          const imgInts = new Int32Array(spliced.length);
+          for (let i = 0; i < spliced.length; i++)
+            imgInts[i] = spliced[i] === vIds!.imageTokenId ? 1 : 0;
+          const imgI32 = MlxArray.fromInt32(imgInts, [spliced.length]);
+          bidirMask = imgI32.astype(Dtype.bool);
+          imgI32.dispose();
+        }
 
-  return { ids: spliced, embeddings: embeds, bidirMask, multimodalMask };
+        return { ids: spliced, embeddings: embeds, bidirMask, multimodalMask, prefixIdentity };
+      } catch (error) { embeds.dispose(); throw error; }
+    });
+  } finally { for (const tensor of reused) tensor?.dispose(); }
 }
 
 /** Vision-only prompt (back-compat wrapper over buildMultimodalPrompt —
@@ -360,14 +394,16 @@ export async function buildVisionPrompt<P extends { softTokens: number }>(
   images: Uint8Array[],
   tokenIds: VisionTokenIds,
   tools: ToolDefinition[] | null = null,
+  nativeWork: PromptNativeWork = (work) => work(),
+  cache?: EncoderCache,
 ): Promise<VisionPrompt> {
   const mp = await buildMultimodalPrompt(
-    model, { vision: { tower, tokenIds } }, tokenizer, template,
-    messages, images, [], tools,
+    model, { vision: { tower, tokenIds, cache } }, tokenizer, template,
+    messages, images, [], tools, nativeWork,
   );
   // no audio was passed, so the union mask IS the image mask — reuse it for
   // the zero-image edge (the original builder returned an all-false mask).
   const imageMask = mp.bidirMask ?? mp.multimodalMask;
   if (mp.bidirMask) mp.multimodalMask.dispose();
-  return { ids: mp.ids, embeddings: mp.embeddings, imageMask };
+  return { ids: mp.ids, embeddings: mp.embeddings, imageMask, prefixIdentity: mp.prefixIdentity };
 }