@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/paged/codec.ts

@@ -0,0 +1,36 @@
+import type { CacheCodec } from "../persistence-types";
+import type { KVCache } from "../kv";
+import { turboQuantFusedDecode } from "../turboquant-codec";
+import { BlockPool, PagedKVCache } from "./cache";
+
+/** Page layout owns its persistence codec, independently of SSD transport.
+ * Affine pages record their bit width and group size, TurboQuant pages their
+ * key and value bit widths; a restored TurboQuant page reads with the read the
+ * restoring graph composed into its cache (`LoadContext.prototype`). */
+export const pagedCacheCodec: CacheCodec = {
+    matches: cache => cache instanceof PagedKVCache,
+    snapshot(cache, context) {
+      const c = cache as PagedKVCache, p = c.pool!;
+      for (const array of p.arrays()) context.push(array, false);
+      return { kind: "paged", offset: c.offset, tensors: context.slots,
+        ...(p.quantization ? p.quantization : {}),
+        ...(p.turbo ? { kBits: p.turbo.kBits, vBits: p.turbo.vBits } : {}),
+        paged: { capacityTokens: c.capacityTokens, blockSize: c.blockSize, blockTable: [...c.blockTable],
+          numBlocks: p.numBlocks, headDim: p.headDim, vHeadDim: p.vHeadDim, dtype: p.dtype, direct: c.direct } };
+    },
+    clone: cache => (cache as PagedKVCache).clone(),
+    load(entry, context) {
+      const layout = entry.paged!;
+      const quantization = entry.bits ? { bits: entry.bits, groupSize: entry.groupSize! } : undefined;
+      const turbo = entry.kBits ? { kBits: entry.kBits, vBits: entry.vBits!, fusedDecode: turboQuantFusedDecode(),
+        kernels: (context.prototype as PagedKVCache).turbo?.kernels ?? (context.prototype as KVCache).turboQuantKernels } : undefined;
+      const cache = new PagedKVCache(layout.capacityTokens, layout.blockSize, layout.direct, quantization, undefined, turbo);
+      const pool = new BlockPool({ ...layout, numKvHeads: entry.tensors[0]!.shape[1]!, quantization,
+        ...(turbo ? { turbo: { kBits: turbo.kBits, vBits: turbo.vBits, fusedDecode: turbo.fusedDecode } } : {}) });
+      try { pool.restoreArrays(entry.tensors.map(context.arr), layout.blockTable); }
+      catch (error) { pool.dispose(); throw error; }
+      cache.pool = pool; cache.offset = entry.offset; cache.blockTable = [...layout.blockTable];
+      return cache;
+    },
+    headerTrimmable: () => true,
+  };

package/src/state/paged/rows.ts

@@ -0,0 +1,223 @@
+import type { AttentionRead, Cache, CommittedAttentionCache, KvAttentionView, KvDonorRows, Mask, SpeculativeRowLayout } from "../../contracts/mlx/cache";
+import { buildBatchedDecodeMask, mergeKVRows } from "../batched-mask";
+import * as ops from "@mlx-bun/mlx/ops";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { disposeResources } from "../../runtime/resources";
+import { PagedKVCache, poolBlocksFor, type PagedQuantization, type PagedTurboQuant } from "./cache";
+
+/** Block storage owns row membership; attention consumes the same gathered
+ * K/V representation as the existing paged cache. No scheduler policy here.
+ *
+ * Reads: fixed at construction by the pages' representation, as the graphs
+ * read these rows before the cache owned the read. Gathered bf16 pages (not
+ * `direct`, no `quantization`): each row appends to its own blocks, the rows'
+ * gathered keys and values join left-padded along the batch, and one fused
+ * SDPA attends under the padding mask of the append (`buildBatchedDecodeMask`),
+ * for decode and windows alike. Reading each row alone over its own length
+ * reduces in another order than this read over the padded length and differs
+ * in the last bits (up to 4e-3 on decode at head dim 512). Direct, affine or
+ * TurboQuant pages: each row reads with its own `PagedKVCache` read for the
+ * phase and the outputs join along the batch. A committed span is the window read, as on
+ * the serial pages.
+ *
+ * The same layout serves continuous decoding, prefill cohorts and
+ * verification: every row owns its blocks and its offset, so the speculative
+ * transaction keeps each row's accepted prefix by trimming that row's blocks
+ * back to it: the rejected tail's blocks return to the row's pool, and the next
+ * appends take the same blocks again. The donor rows are each row's gathered positions,
+ * joined left-padded as the joined read joins them. */
+export class PagedKvRows implements CommittedAttentionCache, SpeculativeRowLayout {
+  #rows: PagedKVCache[] = [];
+  #rope?: MlxArray;
+  #ropeOffset = -1;
+  #beforeRound?: number[];
+  readonly #decode: (keys: MlxArray, values: MlxArray) => AttentionRead;
+  readonly #window: (keys: MlxArray, values: MlxArray) => AttentionRead;
+
+  constructor(readonly capacityTokens: number, readonly blockSize: number, readonly direct = false,
+    readonly quantization?: PagedQuantization, readonly turbo?: PagedTurboQuant) {
+    const alone = direct || quantization !== undefined || turbo !== undefined;
+    this.#decode = alone ? (keys, values) => this.#rowReads(keys, values, (row, k, v) => row.appendDecode(k, v))
+      : (keys, values) => this.#joinedRead(keys, values);
+    this.#window = alone ? (keys, values) => this.#rowReads(keys, values, (row, k, v) => row.appendWindow(k, v))
+      : (keys, values) => this.#joinedRead(keys, values);
+  }
+  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, this.turbo); }
+
+  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)));
+  }
+
+  appendDecode(keys: MlxArray, values: MlxArray): AttentionRead {
+    return this.#decode(keys, values);
+  }
+
+  appendWindow(keys: MlxArray, values: MlxArray): AttentionRead {
+    return this.#window(keys, values);
+  }
+
+  /** The window read, as on the serial pages. */
+  appendCommitted(keys: MlxArray, values: MlxArray): AttentionRead {
+    return this.#window(keys, values);
+  }
+
+  specRoundBegin(): void { this.#beforeRound = this.rowOffsets; }
+  specRoundCommit(): void { this.#beforeRound = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    const before = this.#beforeRound!;
+    for (const [row, storage] of this.#rows.entries())
+      storage.trim(storage.offset - before[row]! - (typeof keep === "number" ? keep : keep[row]!));
+    this.#beforeRound = undefined;
+    this.releaseRopeArr();
+  }
+
+  /** Each row's gathered positions, joined left-padded to the longest row. */
+  captureDonorRows(): KvDonorRows {
+    const gathered: MlxArray[] = [];
+    try {
+      const rows = this.#rows.map(row => {
+        const { keys, values } = row.captureDonorRows();
+        gathered.push(keys, values);
+        return { keys, values };
+      });
+      const { keys, values, leftPad, width } = mergeKVRows(rows);
+      return { keys, values, offsets: this.rowOffsets, starts: leftPad, ends: leftPad.map(() => width) };
+    } finally { disposeResources(gathered); }
+  }
+
+  /** The graphs' read of gathered bf16 rows before the cache owned it: the
+   * padding mask of an `N`-position append built before it, each row's
+   * gathered keys and values joined left-padded along the batch, one fused
+   * SDPA. */
+  #joinedRead(keys: MlxArray, values: MlxArray): AttentionRead {
+    const N = keys.shape[2]!;
+    const mask = buildBatchedDecodeMask(this.batchSize, N, this.offset + N, this.leftPad, null);
+    let joined: [MlxArray, MlxArray];
+    try { joined = this.#joinRows(keys, values); } catch (error) { mask.dispose(); throw error; }
+    const [k, v] = joined;
+    return {
+      attend: (q, scale) => ops.sdpa(q, k, v, scale, "array", mask),
+      dispose() { k.dispose(); v.dispose(); mask.dispose(); },
+    };
+  }
+
+  /** Append each row's slice to its blocks, then join the rows' gathered keys
+   * and values left-padded to the longest row; the caller owns both. */
+  #joinRows(keys: MlxArray, values: MlxArray): [MlxArray, MlxArray] {
+    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.appendAndGather(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])); }
+  }
+
+  #rowReads(keys: MlxArray, values: MlxArray,
+    read: (row: PagedKVCache, k: MlxArray, v: MlxArray) => AttentionRead): AttentionRead {
+    const reads: AttentionRead[] = [];
+    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)]);
+        reads.push(read(storage, k, v));
+      }
+    } catch (error) { disposeResources(reads); throw error; }
+    return {
+      attend(q, scale) {
+        const outputs: MlxArray[] = [];
+        try {
+          for (const [row, rowRead] of reads.entries()) {
+            using query = q.slice([row, 0, 0, 0], [row + 1, ...q.shape.slice(1)]);
+            outputs.push(rowRead.attend(query, scale));
+          }
+          return ops.concatAxis(outputs, 0);
+        } finally { disposeResources(outputs); }
+      },
+      dispose() { disposeResources(reads); },
+    };
+  }
+
+  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);
+    return this.#joinRows(keys, 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.#beforeRound = undefined; this.releaseRopeArr(); disposeResources(rows); }
+}

package/src/state/persistence-types.ts

@@ -0,0 +1,173 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import type { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { type Cache } from "../contracts/mlx/cache";
+import { type CheckpointAttachment } from "../contracts/mlx/checkpoint";
+
+export type CacheKind =
+  | "kv"
+  | "paged"
+  | "rotating"
+  | "qkv"
+  | "rotating-qkv"
+  | "ssm"
+  | "turboquant"
+  | "mla"
+  | "mla-dsa"
+  | "mtp-mla";
+
+/** Header kinds of the latent attention legos (`state/latent-legos.ts`), which
+ * their own codecs (`LATENT_LEGO_CODECS`) save and restore. */
+export type LatentLegoKind = "latent-encoded" | "latent-paged";
+
+export interface TensorSlot {
+  blocks?: Array<{ hash: string; bytes: number }>;
+  off: number;
+  bytes: number;
+  shape: number[];
+  dtype: number;
+  hash: string;
+}
+
+export interface CacheHeaderEntry {
+  paged?: { capacityTokens: number; blockSize: number; blockTable: number[]; numBlocks: number;
+    headDim: number; vHeadDim: number; dtype: Dtype; direct: boolean };
+  kind: CacheKind | LatentLegoKind;
+  offset: number;
+  minimumReusableOffset?: number;
+  /** rotating variants: ring write index */
+  idx?: number;
+  /** rotating variants: window size */
+  maxSize?: number;
+  /** quantized variants (qkv/rotating-qkv: mlx affine scheme) */
+  groupSize?: number;
+  bits?: number;
+  /** turboquant: per-side bit widths (reuses the `bits` field's slot class
+   *  but needs both — asymmetric key/value bit widths, unlike qkv's single
+   *  `bits`). `groupSize` is unused for turboquant (fixed at 32, the
+   *  BLOCK_SIZE constant — not a configurable field like qkv's). */
+  kBits?: number;
+  vBits?: number;
+  /** turboquant: head_dim, needed to unpack kIdx/vPacked on restore
+   *  (packed byte width alone doesn't recover the original element count). */
+  headDim?: number;
+  /** GLM checkpoint-native compressed-cache geometry. */
+  kvLoraRank?: number;
+  ropeHeadDim?: number;
+  dsaHeadDim?: number;
+  maxTokens?: number;
+  /** latent-encoded: the conversion start, and whether the saved row still
+   *  reads f32 (its f32 values follow its encoded planes). */
+  quantizedKvStart?: number;
+  latentShadow?: boolean;
+  /** latent-paged: the row's page layout. */
+  latentPages?: { blockSize: number; blockTable: number[]; numBlocks: number };
+  /** kv/rotating: [k, v] · qkv/rotating-qkv: [kPacked, kScales, kBiases,
+   *  vPacked, vScales, vBiases] · ssm: [conv, recurrent] · turboquant:
+   *  [kIdx, kScales, kZeros, vPacked, vScales] */
+  tensors: TensorSlot[];
+}
+
+export interface KvSaveMeta {
+  attachments?: CheckpointAttachment[];
+  modelId?: string;
+  /** Caller-supplied store identity, including the model and cache scheme. */
+  configFingerprint?: string;
+  /** Adapter namespace (PromptCache ns — adapters joined with "+"). */
+  ns?: string;
+  /** sha256 of tokenizer.json: identical ids must mean identical text. */
+  tokenizerHash?: string;
+  /** In-flight generation state. The serialized caches cover `tokens`;
+   *  pendingToken was sampled from that state but has not been emitted. */
+  generationCheckpoint?: {
+    key: string;
+    cacheNs: string;
+    originalPromptTokens: number;
+    generatedTokens: number;
+    pendingToken: number;
+    /** Actual sampler seed, including a server-generated default. Required to
+     *  continue the random stream exactly after a restart. */
+    seed: number;
+    /** True when `seed` was part of the client request rather than generated
+     *  by the server. */
+    seedWasExplicit: boolean;
+  };
+}
+
+export interface KvFileHeader extends Omit<KvSaveMeta, "attachments"> {
+  attachments?: Array<Omit<CheckpointAttachment, "tensors"> & { tensors: TensorSlot[] }>;
+  /** Absent in earlier v3 files, which used mlx-cache-v3. */
+  codecProvider?: string;
+  formatVersion: 3 | 4 | 5;
+  createdAt: number;
+  tokens: number[];
+  caches: CacheHeaderEntry[];
+}
+
+export interface SnapshotContext {
+  slots: TensorSlot[];
+  push(a: MlxArray, disposeAfter: boolean): void;
+  liveSlice(a: MlxArray, upTo: number): MlxArray;
+  liveMlaSlice(a: MlxArray, upTo: number): MlxArray;
+  pushTriple(t: ops.QuantizedTensor, upTo: number | null): void;
+}
+
+export interface CloneContext {
+  view(a: MlxArray): MlxArray;
+  liveView(a: MlxArray, upTo: number): MlxArray;
+  mlaView(a: MlxArray, upTo: number): MlxArray;
+  tripleView(t: ops.QuantizedTensor, upTo: number | null): ops.QuantizedTensor;
+}
+
+export interface LoadContext {
+  path: string;
+  /** The restoring graph's own empty cache for this entry's layer (from its
+   * `makeCache()`), borrowed for the load: what the graph composed into its
+   * caches (the TurboQuant read a `KVCache` carries) reaches the restored one. */
+  prototype: Cache;
+  arr(slot: TensorSlot): MlxArray;
+  grownArr(slot: TensorSlot, offset: number): MlxArray;
+  triple(slots: TensorSlot[], at: number): ops.QuantizedTensor;
+}
+
+export interface CacheCodec {
+  matches(cache: Cache): boolean;
+  snapshot(cache: Cache, context: SnapshotContext): CacheHeaderEntry;
+  clone(cache: Cache, context: CloneContext): Cache;
+  load(entry: CacheHeaderEntry, context: LoadContext): Cache;
+  headerTrimmable(entry: CacheHeaderEntry): boolean;
+}
+
+/** Codecs belong to the loaded backend binding. The persisted ID is checked
+ * before allocating state; codecs are engine code, never loaded from a file. */
+export interface CacheCodecProvider {
+  readonly id: string;
+  forCache(cache: Cache): CacheCodec;
+  forHeader(entry: CacheHeaderEntry): CacheCodec;
+}
+
+/** Queued immutable snapshot; the queue releases its shared views on settle. */
+export interface SpillItem {
+  attachments?: CheckpointAttachment[];
+  tokens: number[];
+  caches: Cache[];
+  ns: string;
+}
+
+export interface LoadedKvCache {
+  attachments?: CheckpointAttachment[];
+  tokens: number[];
+  header: KvFileHeader;
+  caches: Cache[];
+}
+
+export interface KvLoadExpect {
+  /** Reject on metadata mismatch (pass what the server is running). */
+  modelId?: string;
+  configFingerprint?: string;
+  tokenizerHash?: string;
+  ns?: string;
+  /** Verify every tensor hash before copying it in (off by default,
+   *  `--ssd-cache-verify` — the hash pass roughly doubles restore reads). */
+  verify?: boolean;
+}

package/src/state/persistence-worker.ts

@@ -0,0 +1,83 @@
+import { Worker } from "node:worker_threads";
+import workerSource from "./persistence.worker.js" with { type: "text" };
+import type { KvFileHeader } from "./persistence-types";
+
+/** Immutable, evaluated host-visible storage. The owner retains the native
+ * buffers until write() settles. No MLX handles cross the worker boundary. */
+export interface StoredTensorView {
+  pointer: number;
+  shape: number[];
+  strides: number[];
+  itemSize: number;
+}
+
+export interface KvWriteRequest {
+  layout?: "whole" | "blocks";
+  blockBytes?: number;
+  segmented?: boolean;
+  path: string;
+  header: KvFileHeader;
+  tensors: StoredTensorView[];
+}
+
+/** Hashing, layout packing and atomic disk writes run on a dedicated CPU
+ * worker. Only metadata and pointers are posted; tensor bytes are not copied
+ * through structured clone. The RAM cache remains an independent owner. */
+export class KvWriter {
+  #worker: Worker | null = null;
+  #next = 0;
+  #pending = new Map<number, { resolve(): void; reject(error: Error): void }>();
+
+  lastMetrics: { elapsedMs: number; writtenBytes?: number; reusedBytes?: number; copiedBytes?: number; scratchPeak?: number } | null = null;
+
+  write(request: KvWriteRequest): Promise<void> { return this.#send(request); }
+
+  read(request: { path: string; header: KvFileHeader & { dataStart: number }; pointers: number[]; verify?: boolean }): Promise<void> {
+    return this.#send({ ...request, operation: "read" });
+  }
+
+  collect(root: string): Promise<void> { return this.#send({ root, operation: "collect" }); }
+
+  #send(request: object): Promise<void> {
+    const worker = this.#worker ?? this.#start();
+    worker.ref();
+    const id = ++this.#next;
+    return new Promise<void>((resolve, reject) => {
+      this.#pending.set(id, { resolve, reject });
+      try { worker.postMessage({ id, ...request }); }
+      catch (error) {
+        this.#pending.delete(id);
+        if (!this.#pending.size) worker.unref();
+        reject(error);
+      }
+    });
+  }
+
+  #start(): Worker {
+    // Embed the worker in the single-file executable as well as source runs.
+    const worker = new Worker(workerSource, { eval: true });
+    worker.on("message", (message: { id: number; error?: string; metrics?: KvWriter["lastMetrics"] }) => {
+      const result = this.#pending.get(message.id);
+      if (!result) return;
+      this.#pending.delete(message.id);
+      if (message.error) result.reject(new Error(message.error));
+      else { this.lastMetrics = message.metrics ?? null; result.resolve(); }
+      if (!this.#pending.size) worker.unref();
+    });
+    // Job failures are acknowledged by the worker. A worker-level failure
+    // settles only at exit, after it can no longer read borrowed buffers.
+    worker.on("error", () => {});
+    worker.on("exit", () => {
+      for (const job of this.#pending.values()) job.reject(new Error("KV persistence worker exited"));
+      this.#pending.clear();
+      this.#worker = null;
+    });
+    this.#worker = worker;
+    return worker;
+  }
+}
+
+export const kvWriter = new KvWriter();
+
+/** Reads have their own queue: a restore never waits behind persistence. */
+export const kvReader = new KvWriter();