mlx-bun 0.4.0 → 0.5.0

package/src/lab/paged-kv/paged-attention.ts

@@ -0,0 +1,117 @@
+import { MlxArray } from "../../mlx/array";
+import { Dtype } from "../../mlx/ffi";
+import { MetalKernel } from "../../mlx/metal-kernel";
+import * as ops from "../../mlx/ops";
+import type { KvAttentionView, Mask } from "../../model/gemma4-base";
+import { disposeResources } from "../../engine/resources";
+import type { BlockPool } from "./paged-kv";
+
+// Each SIMD group reduces one query/head over a bounded sequence partition.
+// Only partial output vectors and normalizers materialize, never gathered KV.
+const partialSpec = { name: "paged_attention_partial", ensureRowContiguous: true,
+  inputNames: ["q", "keys", "values", "table", "scale"], outputNames: ["out", "normal"], source: `
+  uint lane = thread_position_in_threadgroup.x;
+  uint hq = threadgroup_position_in_grid.y % HQ;
+  uint query = threadgroup_position_in_grid.y / HQ;
+  uint part = threadgroup_position_in_grid.z;
+  uint hk = hq / (HQ / HK);
+  float accum[(DV + 31) / 32];
+  for (uint j = 0; j < (DV + 31) / 32; ++j) accum[j] = 0;
+  float maximum = -INFINITY, total = 0;
+  uint length = uint(scale[1]);
+  uint count = (length + PART - 1) / PART;
+  uint end = min(uint(length - N + query + 1), (part + 1) * PART);
+  for (uint token = part * PART; token < end; ++token) {
+    uint physical = table[token / BS];
+    uint keyBase = ((physical * HK + hk) * BS + token % BS) * DK;
+    float score = 0;
+    for (uint d = lane; d < DK; d += 32)
+      score += float(q[(hq * N + query) * DK + d]) * float(keys[keyBase + d]);
+    score = simd_sum(score) * scale[0];
+    float next = max(maximum, score);
+    float previous = exp(maximum - next), weight = exp(score - next);
+    uint valueBase = ((physical * HK + hk) * BS + token % BS) * DV;
+    for (uint j = 0; j < (DV + 31) / 32; ++j) {
+      uint d = j * 32 + lane;
+      if (d < DV) accum[j] = accum[j] * previous + weight * float(values[valueBase + d]);
+    }
+    total = total * previous + weight; maximum = next;
+  }
+  uint base = ((hq * N + query) * count + part);
+  for (uint j = 0; j < (DV + 31) / 32; ++j) {
+    uint d = j * 32 + lane;
+    if (d < DV) out[base * DV + d] = accum[j];
+  }
+  if (lane == 0) { normal[base * 2] = maximum; normal[base * 2 + 1] = total; }
+` };
+const partial = new MetalKernel(partialSpec);
+const quantizedPartial = new MetalKernel({ ...partialSpec, name: "paged_attention_affine_partial",
+  inputNames: [...partialSpec.inputNames, "ks", "kb", "vs", "vb"],
+  source: partialSpec.source
+    .replace("float(keys[keyBase + d])", "float(InT(float(ks[(keyBase+d)/GROUP]) * float((keys[(keyBase+d)/PACK] >> (((keyBase+d)%PACK)*BITS)) & ((1u<<BITS)-1u)) + float(kb[(keyBase+d)/GROUP])))")
+    .replace("float(values[valueBase + d])", "float(InT(float(vs[(valueBase+d)/GROUP]) * float((values[(valueBase+d)/PACK] >> (((valueBase+d)%PACK)*BITS)) & ((1u<<BITS)-1u)) + float(vb[(valueBase+d)/GROUP])))"),
+});
+const merge = new MetalKernel({ name: "paged_attention_merge", inputNames: ["parts", "normal"], outputNames: ["out"], source: `
+  uint d = thread_position_in_grid.x;
+  uint row = thread_position_in_grid.y;
+  if (d >= DV) return;
+  uint count = normal_shape[2];
+  float maximum = -INFINITY;
+  for (uint part = 0; part < count; ++part) maximum = max(maximum, normal[(row * count + part) * 2]);
+  float sum = 0, value = 0;
+  for (uint part = 0; part < count; ++part) {
+    uint base = row * count + part;
+    float weight = exp(normal[base * 2] - maximum);
+    sum += normal[base * 2 + 1] * weight;
+    value += parts[base * DV + d] * weight;
+  }
+  out[row * DV + d] = OutT(value / sum);
+` });
+
+/** Independent immutable view: later appends or row retirement cannot alter it. */
+export function pagedAttentionView(pool: BlockPool, blockTable: number[], length: number,
+  direct: boolean): KvAttentionView {
+  const keys = ops.contiguous(pool.keys), values = ops.contiguous(pool.values);
+  const table = MlxArray.fromInt32(Int32Array.from(blockTable), [blockTable.length]);
+  const bs = pool.blockSize, quantization = pool.quantization;
+  const extra = quantization ? [pool.keyScales!, pool.keyBiases!, pool.valueScales!, pool.valueBiases!].map(a => ops.contiguous(a)) : [];
+  const dkStored = pool.headDim, dvStored = pool.vHeadDim;
+  const dtype = pool.dtype;
+  return {
+    attend(q: MlxArray, scale: number, mask: Mask): MlxArray {
+      const [, hq, n, dk] = q.shape as [number, number, number, number];
+      const hk = keys.shape[1]!, dv = dvStored;
+      // Long prefill and explicit masks use the established attention oracle.
+      // This is shape dispatch within one attention interface, not scheduling.
+      if (!direct || n > 8 || mask.mode === "array") {
+        const gather = (tensor: MlxArray) => {
+          using picked = ops.takeAxis(tensor, table, 0);
+          using transposed = ops.transposeAxes(picked, [1, 0, 2, 3]);
+          using flat = ops.reshape(transposed, [1, hk, blockTable.length * bs, tensor.shape[3]!]);
+          return flat.slice([0, 0, 0, 0], [1, hk, length, tensor.shape[3]!]);
+        };
+        using k = gather(keys), v = gather(values);
+        if (!quantization) return ops.sdpa(q, k, v, scale, mask.mode, mask.arr);
+        using ks = gather(extra[0]!), kb = gather(extra[1]!), vs = gather(extra[2]!), vb = gather(extra[3]!);
+        using kd = ops.dequantize(k, ks, kb, { ...quantization, mode: "affine" });
+        using vd = ops.dequantize(v, vs, vb, { ...quantization, mode: "affine" });
+        return ops.sdpa(q, kd, vd, scale, mask.mode, mask.arr);
+      }
+      const part = 64, parts = Math.ceil(length / part);
+      using multiplier = MlxArray.fromFloat32(new Float32Array([scale, length]), [2]);
+      const [out, normal] = (quantization ? quantizedPartial : partial).apply([q, keys, values, table, multiplier, ...extra], {
+        outputs: [{ shape: [hq, n, parts, dv], dtype: Dtype.float32 }, { shape: [hq, n, parts, 2], dtype: Dtype.float32 }],
+        grid: [32, hq * n, parts], threadGroup: [32, 1, 1],
+        templateDtypes: quantization ? { InT: dtype } : {},
+        templateInts: { HQ: hq, HK: hk, N: n, DK: dkStored, DV: dv, BS: bs, PART: part,
+          ...(quantization ? { BITS: quantization.bits, GROUP: quantization.groupSize, PACK: 32 / quantization.bits } : {}) },
+      });
+      try {
+        return merge.apply([out!, normal!], { outputs: [{ shape: [1, hq, n, dv], dtype: q.dtype }],
+          grid: [Math.ceil(dv / 32) * 32, hq * n, 1], threadGroup: [32, 1, 1],
+          templateDtypes: { OutT: q.dtype }, templateInts: { DV: dv } })[0]!;
+      } finally { disposeResources([out!, normal!]); }
+    },
+    dispose() { disposeResources([keys, values, table, ...extra]); },
+  };
+}

package/src/lab/paged-kv/paged-kv-rows.ts

@@ -1,8 +1,10 @@
-import type { BatchableCache, Cache, Mask } from "../../model/gemma4-base";
+import type { BatchableCache, Cache, Mask, KvAttentionView } from "../../model/gemma4-base";
 import { buildBatchedDecodeMask, mergeKVRows } from "../../model/batched-mask";
+import * as ops from "../../mlx/ops";
+import { runtimeFlag } from "../../runtime-config";
 import { MlxArray } from "../../mlx/array";
 import { disposeResources } from "../../engine/resources";
-import { PagedKVCache, poolBlocksFor } from "./paged-kv";
+import { PagedKVCache, poolBlocksFor, type PagedQuantization } from "./paged-kv";
 
 /** Block storage owns row membership; attention consumes the same gathered
  * K/V representation as the existing paged cache. No scheduler policy here. */
@@ -11,7 +13,7 @@ export class PagedKvRows implements BatchableCache {
   #rope?: MlxArray;
   #ropeOffset = -1;
 
-  constructor(readonly capacityTokens: number, readonly blockSize: number) {}
+  constructor(readonly capacityTokens: number, readonly blockSize: number, readonly direct = runtimeFlag("MLX_BUN_PAGED_ATTN", false), readonly quantization?: PagedQuantization) {}
   signature(): string { return "kv:paged-rows"; }
   get batchSize(): number { return this.#rows.length; }
   get offset(): number { return Math.max(0, ...this.#rows.map(row => row.offset)); }
@@ -27,7 +29,7 @@ export class PagedKvRows implements BatchableCache {
     return this.#rope;
   }
   releaseRopeArr(): void { this.#rope?.dispose(); this.#rope = undefined; this.#ropeOffset = -1; }
-  makeEmptyBatch(): PagedKvRows { return new PagedKvRows(this.capacityTokens, this.blockSize); }
+  makeEmptyBatch(): PagedKvRows { return new PagedKvRows(this.capacityTokens, this.blockSize, this.direct, this.quantization); }
 
   mergeRows(rows: readonly Cache[]): void {
     const next: PagedKVCache[] = [];
@@ -51,6 +53,31 @@ export class PagedKvRows implements BatchableCache {
     disposeResources(previous.filter((_, row) => !retained.has(row)));
   }
 
+  get attentionState(): this | undefined { return this.direct || this.quantization ? this : undefined; }
+  appendAndFetch(keys: MlxArray, values: MlxArray): KvAttentionView {
+    const views: KvAttentionView[] = [];
+    try {
+      for (const [row, storage] of this.#rows.entries()) {
+        using k = keys.slice([row, 0, 0, 0], [row + 1, ...keys.shape.slice(1)]);
+        using v = values.slice([row, 0, 0, 0], [row + 1, ...values.shape.slice(1)]);
+        views.push(storage.appendAndFetch(k, v));
+      }
+    } catch (error) { disposeResources(views); throw error; }
+    return {
+      attend(q, scale, mask) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, view] of views.entries()) {
+            using query = q.slice([row, 0, 0, 0], [row + 1, ...q.shape.slice(1)]);
+            outputs.push(view.attend(query, scale, mask));
+          }
+          return ops.concatAxis(outputs, 0);
+        } finally { disposeResources(outputs); }
+      },
+      dispose() { disposeResources(views); },
+    };
+  }
+
   updateAndFetch(keys: MlxArray, values: MlxArray): [MlxArray, MlxArray] {
     if (this.#rows.length === 1) return this.#rows[0]!.updateAndFetch(keys, values);
     const fetched: Array<{ keys: MlxArray; values: MlxArray }> = [];
@@ -66,6 +93,7 @@ export class PagedKvRows implements BatchableCache {
     } finally { disposeResources(fetched.flatMap(row => [row.keys, row.values])); }
   }
   makeMask(tokens: number, window: number | null): Mask {
+    if (this.attentionState && window === null) return { mode: tokens === 1 ? "" : "causal", arr: null };
     if (this.#rows.length === 1) return this.#rows[0]!.makeMask(tokens, window);
     return { mode: "array", arr: buildBatchedDecodeMask(this.batchSize, tokens,
       this.offset + tokens, this.leftPad, window) };
@@ -76,7 +104,7 @@ export class PagedKvRows implements BatchableCache {
   projectedBytes(tokens: number): number {
     const pool = this.#rows.find(row => row.pool)?.pool;
     if (!pool) return 0;
-    const bytesPerBlock = (pool.keys.nbytes + pool.values.nbytes) / pool.numBlocks;
+    const bytesPerBlock = pool.arrays().reduce((bytes, array) => bytes + array.nbytes, 0) / pool.numBlocks;
     return poolBlocksFor(tokens, this.blockSize) * bytesPerBlock;
   }
   dispose(): void { const rows = this.#rows; this.#rows = []; this.releaseRopeArr(); disposeResources(rows); }

package/src/lab/paged-kv/paged-kv.ts

@@ -1,29 +1,16 @@
-// Paged KV cache — vLLM-style block-pool storage behind the standard Cache
-// interface (docs/design/kv-cache.md). OPTIONAL and default-off
-// (`--paged-kv` / GenerateOptions.pagedKv); Gemma4-family plain full-attention
-// layers, bf16 only. PagedKvRows composes independent pools at B1/B>1.
-//
-// Storage: K/V live in fixed-size per-layer pool tensors
-// [numBlocks, H_kv, blockSize, headDim]; a host-side block table maps the
-// sequence's logical block order to physical pool slots. Writes go through
-// ops.sliceUpdate into the tail block; reads gather the occupied blocks
-// back into ONE contiguous [1, H, S, D] pair (ops.takeAxis over the pool's
-// block axis + transpose + reshape) and hand that to the unchanged
-// ops.sdpa call site — identical bytes to what a plain KVCache would have
-// fetched, so the paged path is gated bit-exact, not KL-tolerated
-// (tests/parity/paged-kv-parity.test.ts). The gather IS the cost: a full K/V copy
-// per step, pure bandwidth tax at batch=1 — v1 ships the abstraction for
-// the batched/CoW follow-ups, not a speed win (see the design doc's
-// Motivation, which says this honestly).
-//
-// Deliberately NOT a KVCache subclass (the TurboQuantKVCache reasoning):
-// compiled graph adapters exclude the changing block-list geometry. The
-// cache-layout binding supplies dynamic rows independently of scheduling.
+// Paged KV storage behind the common cache and attention interfaces.
+// Full-attention pools support bf16 and affine KV4/KV8. The established
+// gather + SDPA arm is the numerical control; direct Metal reads are Lab.
+// Published views and SSD snapshots own immutable buffers independently of
+// appends, row membership and RAM eviction.
 
+import { disposeResources } from "../../engine/resources";
+import { pagedAttentionView } from "./paged-attention";
+import { runtimeFlag } from "../../runtime-config";
 import { MlxArray } from "../../mlx/array";
 import { Dtype } from "../../mlx/ffi";
 import * as ops from "../../mlx/ops";
-import { createCausalMask, type Cache, type Mask } from "../../model/gemma4-base";
+import { createCausalMask, type Cache, type Mask, type KvAttentionView } from "../../model/gemma4-base";
 
 /** Typed pool-exhaustion error: a generation outgrew its pool. Sizing from
  *  prompt+maxTokens at construction makes this unreachable in practice;
@@ -41,10 +28,18 @@ export class PagedPoolExhausted extends Error {
  *  (ops.zeros — no host-pointer alignment/dtor hazards), shaped
  *  [numBlocks, H, blockSize, D]. alloc()/free() manage physical slot
  *  indices; the owning PagedKVCache writes/gathers through the pool. */
+export interface PagedQuantization { bits: 4 | 8; groupSize: number }
+
 export class BlockPool {
+  keyScales?: MlxArray; keyBiases?: MlxArray;
+  valueScales?: MlxArray; valueBiases?: MlxArray;
+  readonly headDim: number;
+  readonly vHeadDim: number;
+  readonly dtype: Dtype;
+  readonly quantization?: PagedQuantization;
   keys: MlxArray;
   values: MlxArray;
-  readonly numBlocks: number;
+  numBlocks: number;
   readonly blockSize: number;
   #free: number[];
 
@@ -57,7 +52,10 @@ export class BlockPool {
      *  reach v1's gemma4 scope, but the pool stays shape-honest). */
     vHeadDim?: number;
     dtype: Dtype;
+    quantization?: PagedQuantization;
   }, source?: BlockPool) {
+    this.headDim = opts.headDim; this.vHeadDim = opts.vHeadDim ?? opts.headDim;
+    this.dtype = opts.dtype; this.quantization = opts.quantization;
     this.numBlocks = opts.numBlocks;
     this.blockSize = opts.blockSize;
     this.keys = source ? ops.contiguous(source.keys) : ops.zeros(
@@ -65,6 +63,16 @@ export class BlockPool {
     this.values = source ? ops.contiguous(source.values) : ops.zeros(
       [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.vHeadDim ?? opts.headDim],
       opts.dtype);
+    if (source?.quantization) {
+      this.keyScales = ops.contiguous(source.keyScales!); this.keyBiases = ops.contiguous(source.keyBiases!);
+      this.valueScales = ops.contiguous(source.valueScales!); this.valueBiases = ops.contiguous(source.valueBiases!);
+    } else if (this.quantization) {
+      const { bits, groupSize } = this.quantization;
+      const k = ops.quantize(this.keys, groupSize, bits), v = ops.quantize(this.values, groupSize, bits);
+      this.keys.dispose(); this.values.dispose();
+      this.keys = k.packed; this.keyScales = k.scales; this.keyBiases = k.biases;
+      this.values = v.packed; this.valueScales = v.scales; this.valueBiases = v.biases;
+    }
     // LIFO free list, low indices first — deterministic layout for tests.
     this.#free = source ? [...source.#free] : Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
   }
@@ -73,8 +81,31 @@ export class BlockPool {
    * Functional writes detach storage while another snapshot retains it. */
   clone(): BlockPool {
     return new BlockPool({ numBlocks: this.numBlocks, blockSize: this.blockSize,
-      numKvHeads: this.keys.shape[1]!, headDim: this.keys.shape[3]!,
-      vHeadDim: this.values.shape[3]!, dtype: this.keys.dtype }, this);
+      numKvHeads: this.keys.shape[1]!, headDim: this.headDim,
+      vHeadDim: this.vHeadDim, dtype: this.dtype, quantization: this.quantization }, this);
+  }
+
+  restoreArrays(arrays: MlxArray[], occupied: readonly number[]): void {
+    disposeResources(this.arrays());
+    [this.keys, this.values] = [arrays[0]!, arrays[1]!];
+    [this.keyScales, this.keyBiases, this.valueScales, this.valueBiases] = arrays.slice(2);
+    const used = new Set(occupied);
+    this.#free = Array.from({ length: this.numBlocks }, (_, i) => this.numBlocks - 1 - i).filter(i => !used.has(i));
+  }
+
+  grow(blocks: number): void {
+    if (blocks <= this.numBlocks) return;
+    const grow = (array: MlxArray) => {
+      using zeros = ops.zeros([blocks - this.numBlocks, ...array.shape.slice(1)], array.dtype);
+      const next = ops.concatAxis([array, zeros], 0); array.dispose(); return next;
+    };
+    this.keys = grow(this.keys); this.values = grow(this.values);
+    if (this.quantization) {
+      this.keyScales = grow(this.keyScales!); this.keyBiases = grow(this.keyBiases!);
+      this.valueScales = grow(this.valueScales!); this.valueBiases = grow(this.valueBiases!);
+    }
+    for (let i = blocks - 1; i >= this.numBlocks; i--) this.#free.push(i);
+    this.numBlocks = blocks;
   }
 
   get freeBlocks(): number {
@@ -96,6 +127,20 @@ export class BlockPool {
    *  the pool tensors are single-referenced here, so mlx donates the
    *  buffer and the write is in-place in the steady state. */
   writeBlock(physIdx: number, k: MlxArray, v: MlxArray, within: number): void {
+    if (this.quantization) {
+      const { bits, groupSize } = this.quantization;
+      const kq = ops.quantize(k, groupSize, bits), vq = ops.quantize(v, groupSize, bits);
+      const replace = (target: MlxArray, input: MlxArray): MlxArray => {
+        const result = ops.sliceUpdate(target, input, [physIdx, 0, within, 0],
+          [physIdx + 1, input.shape[1]!, within + input.shape[2]!, input.shape[3]!]);
+        target.dispose(); return result;
+      };
+      try {
+        this.keys = replace(this.keys, kq.packed); this.keyScales = replace(this.keyScales!, kq.scales); this.keyBiases = replace(this.keyBiases!, kq.biases);
+        this.values = replace(this.values, vq.packed); this.valueScales = replace(this.valueScales!, vq.scales); this.valueBiases = replace(this.valueBiases!, vq.biases);
+      } finally { disposeResources([kq.packed, kq.scales, kq.biases, vq.packed, vq.scales, vq.biases]); }
+      return;
+    }
     const l = k.shape[2]!;
     const [, H, , kD] = this.keys.shape as [number, number, number, number];
     const vD = this.values.shape[3]!;
@@ -122,13 +167,23 @@ export class BlockPool {
       t.dispose();
       return out;
     };
-    const k = pick(this.keys);
-    const v = pick(this.values);
+    const read = (data: MlxArray, scales?: MlxArray, biases?: MlxArray) => {
+      const packed = pick(data);
+      if (!this.quantization) return packed;
+      using owned = packed; using s = pick(scales!), b = pick(biases!);
+      return ops.dequantize(owned, s, b, { ...this.quantization, mode: "affine" });
+    };
+    const k = read(this.keys, this.keyScales, this.keyBiases);
+    const v = read(this.values, this.valueScales, this.valueBiases);
     idx.dispose();
     return [k, v];
   }
 
+  arrays(): MlxArray[] {
+    return [this.keys, this.values, this.keyScales, this.keyBiases, this.valueScales, this.valueBiases].filter((a): a is MlxArray => !!a);
+  }
   dispose(): void {
+    disposeResources([this.keyScales, this.keyBiases, this.valueScales, this.valueBiases].filter((a): a is MlxArray => !!a));
     this.keys.dispose();
     this.values.dispose();
     this.#free = [];
@@ -159,10 +214,12 @@ export class PagedKVCache implements Cache {
     /** Tokens this cache must be able to hold (prompt + maxTokens). */
     readonly capacityTokens: number,
     readonly blockSize: number,
+    readonly direct = runtimeFlag("MLX_BUN_PAGED_ATTN", false),
+    readonly quantization?: PagedQuantization,
   ) {}
 
   clone(): PagedKVCache {
-    const copy = new PagedKVCache(this.capacityTokens, this.blockSize);
+    const copy = new PagedKVCache(this.capacityTokens, this.blockSize, this.direct, this.quantization);
     copy.offset = this.offset;
     copy.blockTable = [...this.blockTable];
     copy.pool = this.pool?.clone() ?? null;
@@ -178,7 +235,23 @@ export class PagedKVCache implements Cache {
     return Math.ceil(n / this.#blockSize);
   }
 
+  get attentionState(): this | undefined { return this.direct || this.quantization ? this : undefined; }
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView {
+    this.append(k, v);
+    return pagedAttentionView(this.pool!, [...this.blockTable], this.offset, this.direct);
+  }
+
   updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    this.append(k, v);
+    const H = this.pool!.keys.shape[1]!, kD = this.pool!.headDim, vD = this.pool!.vHeadDim;
+    const [gk, gv] = this.pool!.gather(this.blockTable);
+    const keys = gk.slice([0, 0, 0, 0], [1, H, this.offset, kD]);
+    const values = gv.slice([0, 0, 0, 0], [1, H, this.offset, vD]);
+    gk.dispose(); gv.dispose();
+    return [keys, values];
+  }
+
+  append(k: MlxArray, v: MlxArray): void {
     const L = k.shape[2]!;
     const [, H, , kD] = k.shape as [number, number, number, number];
     const vD = v.shape[3]!;
@@ -189,7 +262,7 @@ export class PagedKVCache implements Cache {
         numKvHeads: H,
         headDim: kD,
         vHeadDim: vD,
-        dtype: k.dtype,
+        dtype: k.dtype, quantization: this.quantization,
       });
     // Write, splitting the incoming L along block boundaries. Blocks past
     // the current tail allocate from the free list as they're reached.
@@ -198,7 +271,10 @@ export class PagedKVCache implements Cache {
       const pos = this.offset + written;
       const bi = Math.floor(pos / this.#blockSize);
       const within = pos % this.#blockSize;
-      while (this.blockTable.length <= bi) this.blockTable.push(this.pool.alloc());
+      while (this.blockTable.length <= bi) {
+        if (!this.pool.freeBlocks) this.pool.grow(Math.max(this.pool.numBlocks + 1, Math.ceil(this.pool.numBlocks * 1.5)));
+        this.blockTable.push(this.pool.alloc());
+      }
       const l = Math.min(this.#blockSize - within, L - written);
       // Full-range pieces skip the slice (fresh-view slice would be a
       // gratuitous op); partial pieces slice the [written, written+l) rows.
@@ -215,14 +291,6 @@ export class PagedKVCache implements Cache {
     }
     this.offset += L;
 
-    // Fetch: gather occupied blocks contiguous, slice to the live prefix.
-    const nb = this.#blocksFor(this.offset);
-    const [gk, gv] = this.pool.gather(this.blockTable.slice(0, nb));
-    const keys = gk.slice([0, 0, 0, 0], [1, H, this.offset, kD]);
-    const values = gv.slice([0, 0, 0, 0], [1, H, this.offset, vD]);
-    gk.dispose();
-    gv.dispose();
-    return [keys, values];
   }
 
   /** Same mask policy as KVCache.makeMask (it reads only offset). */
@@ -233,11 +301,9 @@ export class PagedKVCache implements Cache {
     return { mode: "array", arr: createCausalMask(N, this.offset, windowSize) };
   }
 
-  /** Pool tensors — what a prefill-chunk boundary must materialize.
-   *  NOT prompt-cache-compatible in v1 (cloneKvCaches never sees paged
-   *  caches; the serve lane bypasses take/put for paged requests). */
+  /** Owned planes materialized at prefill and immutable publication boundaries. */
   state(): MlxArray[] {
-    return this.pool ? [this.pool.keys, this.pool.values] : [];
+    return this.pool ? this.pool.arrays() : [];
   }
 
   isTrimmable(): boolean {

package/src/md-modules.d.ts

@@ -6,6 +6,11 @@ declare module "*.md" {
   export default content;
 }
 
+declare module "*.worker.js" {
+  const content: string;
+  export default content;
+}
+
 // Same deal for vendored text-imported assets under src/web/vendor/ (see
 // that dir's README) — served as raw strings via `with { type: "text" }`,
 // but bun-types has no ambient declaration for .js/.css read as text.

package/src/mlx/array.ts

@@ -2,6 +2,7 @@
 // Explicit .dispose() is the contract; a FinalizationRegistry backstop
 // frees leaked handles on GC (verified in lab/spikes/phase0-memory.ts).
 
+import { hostBufferDestructor, type HostBuffer } from "../storage/host-buffer";
 import { dlopen, ptr, toArrayBuffer } from "bun:ffi";
 import { C, Dtype, DTYPE_NAMES, type MlxHandle, optInt, outArray, takeMlxError } from "./ffi";
 import type { SafetensorsDtype } from "../safetensors";
@@ -130,6 +131,18 @@ export class MlxArray {
     return new MlxArray(handle);
   }
 
+  /** Adopt a page-aligned CPU restore allocation without a host copy.
+   * The native destructor owns it after this call, including GPU references. */
+  static adoptHostBuffer(buffer: HostBuffer, shape: number[], dtype: Dtype): MlxArray {
+    const sb = shapeBuf(shape);
+    const handle = C.mlx_array_new_data_managed_payload(
+      buffer.pointer, ptr(sb), shape.length, dtype, buffer.pointer, hostBufferDestructor,
+    );
+    const array = new MlxArray(handle);
+    buffer.transfer();
+    return array;
+  }
+
   /** Copying constructor for small host data. */
   static fromFloat32(data: Float32Array, shape: number[]): MlxArray {
     const sb = shapeBuf(shape);
@@ -237,6 +250,30 @@ export class MlxArray {
     return this.rawBytesView().slice();
   }
 
+  /** Publish an evaluated immutable buffer to a CPU reader. Strides are in
+   * elements, so padding, transposes and slices need no GPU packing copy.
+   * The caller retains this array until the reader reports completion. */
+  storageView(): import("../storage/kv-writer").StoredTensorView {
+    this.eval();
+    const shape = this.shape;
+    const stridePtr = C.mlx_array_strides(this.handle);
+    const strides = shape.length
+      ? [...new BigInt64Array(toArrayBuffer(stridePtr!, 0, shape.length * 8))].map(Number)
+      : [];
+    const dt = this.dtype;
+    const pointer =
+      dt === Dtype.float32 ? C.mlx_array_data_float32(this.handle)
+      : dt === Dtype.float16 ? C.mlx_array_data_float16(this.handle)
+      : dt === Dtype.bfloat16 ? C.mlx_array_data_bfloat16(this.handle)
+      : dt === Dtype.uint32 ? C.mlx_array_data_uint32(this.handle)
+      : dt === Dtype.int32 ? C.mlx_array_data_int32(this.handle)
+      : dt === Dtype.uint8 ? C.mlx_array_data_uint8(this.handle)
+      : dt === Dtype.int8 ? C.mlx_array_data_int8(this.handle)
+      : null;
+    if (pointer === null && this.size) throw new Error(`storageView: unsupported or unavailable ${this.dtypeName}`);
+    return { pointer: Number(pointer), shape, strides, itemSize: Number(C.mlx_array_itemsize(this.handle)) };
+  }
+
   /** ZERO-COPY view of the evaluated array's bytes — aliases the mlx
    *  buffer directly, no JS-heap copy. Valid only while THIS array is
    *  alive and unmutated: do not retain past dispose(). Built for

package/src/mlx/ffi.ts

@@ -53,6 +53,8 @@ export const C = dlopen(LIBMLXC_PATH, {
   // array metadata
   mlx_array_ndim: { args: [u64], returns: u64 },
   mlx_array_shape: { args: [u64], returns: P },
+  mlx_array_strides: { args: [u64], returns: P },
+  mlx_array_itemsize: { args: [u64], returns: u64 },
   mlx_array_dim: { args: [u64, i32], returns: i32 },
   mlx_array_dtype: { args: [u64], returns: i32 },
   mlx_array_size: { args: [u64], returns: u64 },
@@ -181,6 +183,23 @@ export const C = dlopen(LIBMLXC_PATH, {
   // (res, input, weight, stride, padding, dilation, groups, stream) — depthwise
   // causal conv1d for Qwen3.5 gated-DeltaNet (groups == channels, padding 0).
   mlx_conv1d: { args: [P, u64, u64, i32, i32, i32, i32, u64], returns: i32 },
+  // --- Whisper (src/model/whisper.ts, src/audio/whisper-mel.ts) -----------
+  // (res, a, n, axis, norm, stream) — mlx_fft_norm enum: 0 backward (numpy default).
+  mlx_fft_rfft: { args: [P, u64, i32, i32, i32, u64], returns: i32 },
+  // (res, a, shape*, shape_num, strides*(int64), strides_num, offset, stream)
+  mlx_as_strided: { args: [P, u64, P, u64, P, u64, u64, u64], returns: i32 },
+  // (res, a, axes*, n, low*, n, high*, n, pad_value, mode, stream)
+  mlx_pad: { args: [P, u64, P, u64, P, u64, P, u64, u64, cstring, u64], returns: i32 },
+  mlx_log10: { args: [P, u64, u64], returns: i32 },
+  // full reduce: (res, a, keepdims, stream)
+  mlx_max: { args: [P, u64, FFIType.bool, u64], returns: i32 },
+  mlx_broadcast_to: { args: [P, u64, P, u64, u64], returns: i32 },
+  mlx_topk_axis: { args: [P, u64, i32, i32, u64], returns: i32 },
+  mlx_var_axis: { args: [P, u64, i32, FFIType.bool, i32, u64], returns: i32 },
+  mlx_stack_axis: { args: [P, u64, i32, u64], returns: i32 },
+  mlx_minimum: { args: [P, u64, u64, u64], returns: i32 },
+  // flat take: (res, a, indices, stream)
+  mlx_take: { args: [P, u64, u64, u64], returns: i32 },
   // Natural mlx-c signatures. Bun 1.4 fixed the macOS arm64 stack packing
   // bug recorded in lab/repro/bun-ffi-stack-args, so these must match the
   // header exactly. mlx-bun now requires Bun 1.4 or newer.

package/src/mlx/normalized-argmax.ts

@@ -0,0 +1,116 @@
+// Greedy selection preserves dtype-rounded normalized scores, including ties.
+// Native logsumexp retains the reference reduction. Tiled selection avoids
+// writing the vocabulary-sized normalized array. Inputs remain borrowed.
+import { MlxArray } from "./array";
+import { Dtype } from "./ffi";
+import * as ops from "./ops";
+import { CompiledFunction, isShapelessTracing } from "./compile";
+import { MetalKernel } from "./metal-kernel";
+
+const header = `
+inline bool better(float x, uint xi, float y, uint yi) {
+  return x > y || (x == y && xi < yi);
+}`;
+
+// Four simdgroups reduce score/index pairs. Lowest token ID wins every tie.
+// The caller finishes the simd==0 block by writing its output from lane zero.
+const reduction = `
+  for (uint distance = 16; distance > 0; distance >>= 1) {
+    float other = simd_shuffle_down(best, distance);
+    uint oi = simd_shuffle_down(bestId, distance);
+    if (lane + distance < 32 && better(other, oi, best, bestId)) {
+      best = other; bestId = oi;
+    }
+  }
+  threadgroup float scores[4];
+  threadgroup uint ids[4];
+  if (lane == 0) { scores[simd] = best; ids[simd] = bestId; }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+  if (simd == 0) {
+    best = lane < 4 ? scores[lane] : -INFINITY;
+    bestId = lane < 4 ? ids[lane] : 0xffffffffu;
+    for (uint distance = 2; distance > 0; distance >>= 1) {
+      float other = simd_shuffle_down(best, distance);
+      uint oi = simd_shuffle_down(bestId, distance);
+      if (lane + distance < 4 && better(other, oi, best, bestId)) {
+        best = other; bestId = oi;
+      }
+    }
+`;
+
+const partialSource = `
+  const uint tid = thread_position_in_threadgroup.x;
+  const uint lane = tid % 32, simd = tid / 32;
+  const uint row = threadgroup_position_in_grid.y;
+  const uint tile = threadgroup_position_in_grid.x;
+  float best = -INFINITY;
+  uint bestId = 0xffffffffu;
+  for (uint i = tile * 1024 + tid; i < min(uint(V), (tile + 1) * 1024); i += 128) {
+    float value = float(T(logits[row * V + i] - lse[row]));
+    // A nonfinite logsumexp produces only NaN/-inf normalized scores.
+    // Native argmax selects token zero for this degenerate row.
+    value = isnan(value) ? -INFINITY : value;
+    if (better(value, i, best, bestId)) { best = value; bestId = i; }
+  }
+  ${reduction}
+    if (lane == 0) { values[row * P + tile] = best; indices[row * P + tile] = bestId; }
+  }
+`;
+
+const finishSource = `
+  const uint tid = thread_position_in_threadgroup.x;
+  const uint lane = tid % 32, simd = tid / 32;
+  const uint row = threadgroup_position_in_grid.y;
+  float best = -INFINITY;
+  uint bestId = 0xffffffffu;
+  for (uint i = tid; i < P; i += 128) {
+    float value = values[row * P + i];
+    uint id = indices[row * P + i];
+    if (better(value, id, best, bestId)) { best = value; bestId = id; }
+  }
+  ${reduction}
+    if (lane == 0) out[row] = bestId;
+  }
+`;
+
+let partial: MetalKernel | undefined;
+let finish: MetalKernel | undefined;
+let compiled: CompiledFunction | undefined;
+
+function kernelGraph(logits: MlxArray): MlxArray {
+  const shape = logits.shape;
+  const vocab = shape.at(-1)!, rows = logits.size / vocab, tiles = Math.ceil(vocab / 1024);
+  partial ??= new MetalKernel({ name: "normalized_argmax_partial", header,
+    inputNames: ["logits", "lse"], outputNames: ["values", "indices"], source: partialSource });
+  finish ??= new MetalKernel({ name: "normalized_argmax_finish", header,
+    inputNames: ["values", "indices"], outputNames: ["out"], source: finishSource });
+  using lse = ops.logsumexpAxis(logits, -1, true);
+  const [values, indices] = partial.apply([logits, lse], {
+    outputs: [{ shape: [rows, tiles], dtype: Dtype.float32 }, { shape: [rows, tiles], dtype: Dtype.uint32 }],
+    grid: [tiles * 128, rows, 1], threadGroup: [128, 1, 1],
+    templateInts: { V: vocab, P: tiles }, templateDtypes: { T: logits.dtype },
+  });
+  try {
+    using tokens = finish.apply([values!, indices!], {
+      outputs: [{ shape: [rows], dtype: Dtype.uint32 }],
+      grid: [128, rows, 1], threadGroup: [128, 1, 1], templateInts: { P: tiles },
+    })[0]!;
+    return ops.reshape(tokens, shape.slice(0, -1));
+  } finally { values!.dispose(); indices!.dispose(); }
+}
+
+/** Return owned greedy IDs with the same shape/rounding as argmax(logprobs).
+ * Fixed-shape compilation caches graph/FFI work for repeated row geometries.
+ * Shapeless enclosing graphs retain native operations, which support symbolic
+ * sizes; integer/empty inputs also retain the original operation semantics. */
+export function normalizedArgmax(logits: MlxArray): MlxArray {
+  const dtype = logits.dtype;
+  if (isShapelessTracing() || logits.size === 0 ||
+      (dtype !== Dtype.float32 && dtype !== Dtype.float16 && dtype !== Dtype.bfloat16)) {
+    using lse = ops.logsumexpAxis(logits, -1, true);
+    using scores = ops.sub(logits, lse);
+    return ops.argmaxAxis(scores, -1);
+  }
+  compiled ??= new CompiledFunction(([input]) => [kernelGraph(input!)], false);
+  return compiled.apply([logits])[0]!;
+}