mlx-bun 0.3.0 → 0.5.0

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

@@ -0,0 +1,111 @@
+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, 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. */
+export class PagedKvRows implements BatchableCache {
+  #rows: PagedKVCache[] = [];
+  #rope?: MlxArray;
+  #ropeOffset = -1;
+
+  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)); }
+  get rowOffsets(): number[] { return this.#rows.map(row => row.offset); }
+  get leftPad(): number[] { const width = this.offset; return this.#rows.map(row => width - row.offset); }
+  get ropeOffsetArr(): MlxArray | undefined {
+    if (this.#rows.length < 2) return undefined;
+    if (!this.#rope || this.#ropeOffset !== this.offset) {
+      this.releaseRopeArr();
+      this.#rope = MlxArray.fromInt32(Int32Array.from(this.rowOffsets), [this.batchSize]);
+      this.#ropeOffset = this.offset;
+    }
+    return this.#rope;
+  }
+  releaseRopeArr(): void { this.#rope?.dispose(); this.#rope = undefined; this.#ropeOffset = -1; }
+  makeEmptyBatch(): PagedKvRows { return new PagedKvRows(this.capacityTokens, this.blockSize, this.direct, this.quantization); }
+
+  mergeRows(rows: readonly Cache[]): void {
+    const next: PagedKVCache[] = [];
+    try {
+      for (const source of rows) {
+        if (source instanceof PagedKvRows) {
+          for (const row of source.#rows) next.push(row.clone());
+        } else if (source instanceof PagedKVCache) next.push(source.clone());
+        else throw new Error(`Paged rows cannot merge ${source.signature()}`);
+      }
+    } catch (error) { disposeResources(next); throw error; }
+    const previous = this.#rows; this.#rows = next;
+    this.releaseRopeArr(); disposeResources(previous);
+  }
+  extractRow(row: number): PagedKVCache { return this.#rows[row]!.clone(); }
+  filterRows(keep: readonly number[]): void {
+    const previous = this.#rows;
+    this.#rows = keep.map(row => previous[row]!);
+    const retained = new Set(keep);
+    this.releaseRopeArr();
+    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 }> = [];
+    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)]);
+        const [nextK, nextV] = storage.updateAndFetch(k, v);
+        fetched.push({ keys: nextK, values: nextV });
+      }
+      const result = mergeKVRows(fetched);
+      return [result.keys, result.values];
+    } 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) };
+  }
+  state(): MlxArray[] { return this.#rows.flatMap(row => row.state()); }
+  isTrimmable(): boolean { return true; }
+  trim(tokens: number): void { for (const row of this.#rows) row.trim(tokens); this.releaseRopeArr(); }
+  projectedBytes(tokens: number): number {
+    const pool = this.#rows.find(row => row.pool)?.pool;
+    if (!pool) return 0;
+    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,30 +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); v1 scope is serial batch=1
-// decode on Gemma4-family plain full-attention layers, bf16 only.
-//
-// 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):
-// every instanceof gate in the tree — CompiledDecode.supports, the batch
-// gateway's #modelCachesBatchable, generated forwards' #matches() — must
-// EXCLUDE paged caches and fall back to the monolith/uncompiled path.
+// 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;
@@ -42,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[];
 
@@ -58,16 +52,60 @@ 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 = ops.zeros(
+    this.keys = source ? ops.contiguous(source.keys) : ops.zeros(
       [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.headDim], opts.dtype);
-    this.values = ops.zeros(
+    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 = Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
+    this.#free = source ? [...source.#free] : Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
+  }
+
+  /** Independent ownership and allocation metadata over immutable MLX buffers.
+   * 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.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 {
@@ -89,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]!;
@@ -115,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 = [];
@@ -134,8 +196,7 @@ export class BlockPool {
  *  shape): head count / head dims / dtype come from the first k/v pair,
  *  so the wiring (maybePageKv) needs no per-model shape plumbing. */
 export class PagedKVCache implements Cache {
-  /** Distinct kind: paged layout is serial-only and never merges into a
-   *  batch, so no capability guard should ever match it as plain. */
+  /** Distinct storage kind; the row adapter preserves its block pools. */
   signature(): string { return "kv:paged"; }
   /** Matches KVCache.STEP: v1's growth granularity is a permutation of
    *  today's 256-token step into fixed reusable slots, not a new tuning
@@ -153,8 +214,18 @@ 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, this.direct, this.quantization);
+    copy.offset = this.offset;
+    copy.blockTable = [...this.blockTable];
+    copy.pool = this.pool?.clone() ?? null;
+    return copy;
+  }
+
   get #blockSize(): number {
     return this.blockSize;
   }
@@ -164,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]!;
@@ -175,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.
@@ -184,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.
@@ -201,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). */
@@ -219,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);
@@ -232,11 +245,35 @@ export class MlxArray {
     );
   }
 
-  /** Raw bytes of the evaluated array (copy). bf16/f16/f32 only. */
+  /** Raw bytes of a supported floating-point or integer array (copy). */
   rawBytes(): Uint8Array {
     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
@@ -253,6 +290,7 @@ export class MlxArray {
       : 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;

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]!;
+}