@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/rows.ts

@@ -0,0 +1,104 @@
+import type { BatchableCache, Cache } from "../contracts/mlx/cache";
+import { applyStateChanges, cleanupFailure, disposeResources } from "../runtime/resources";
+import type { PreparedStateChange } from "../contracts/portable/resources";
+
+/** Preserve the graph's layout capabilities when an empty sibling is made. */
+export interface MlxRowLayout<Layout extends BatchableCache> extends BatchableCache {
+  makeEmptyBatch(): Layout;
+}
+
+/** Layer layouts own tensor geometry. This collection only applies the same
+ * row membership operation to every layer. It has no scheduling, sampling,
+ * checkpoint-key or persistence policy. */
+export class MlxStateRows<Layout extends MlxRowLayout<Layout> = BatchableCache> {
+  /** Borrowed by the bound target graph; collection ownership stays here. */
+  readonly caches: Layout[];
+  #rows = 0;
+
+  /** Takes ownership of empty, compatible layer layouts. */
+  constructor(layouts: readonly Layout[]) { this.caches = [...layouts]; }
+  get rowCount(): number { return this.#rows; }
+  /** Retain an immutable batch boundary without extracting or materializing
+   * individual rows. Layer layouts decide how to share their backing arrays. */
+  clone(): MlxStateRows<Layout> {
+    const layouts: Layout[] = [];
+    try {
+      for (const source of this.caches) {
+        const layout = source.makeEmptyBatch(); layouts.push(layout);
+        layout.mergeRows([source]);
+      }
+      const result = new MlxStateRows(layouts); result.#rows = this.#rows;
+      return result;
+    } catch (error) { return cleanupFailure(error, () => disposeResources(layouts)); }
+  }
+  /** Prepare an initial cohort with one merge per layer. */
+  mergeRows(rows: readonly (readonly Cache[])[]): void {
+    const next: Layout[] = [];
+    try {
+      for (let layer = 0; layer < this.caches.length; layer++) {
+        const layout = this.caches[layer]!.makeEmptyBatch(); next.push(layout);
+        layout.mergeRows(rows.map(row => row[layer]!));
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.caches.splice(0, this.caches.length, ...next);
+    this.#rows = rows.length; disposeResources(previous);
+  }
+
+  /** Borrow a prepared request's state. Build every replacement before
+   * releasing the active group, so a failed admission leaves it intact. */
+  append(row: readonly Cache[]): void {
+    applyStateChanges([() => this.prepareAppend(row)]);
+  }
+
+  /** Stage membership so a method can compose it with companion-state admission. */
+  prepareAppend(row: readonly Cache[]): PreparedStateChange {
+    return this.prepareAppendMany([row]);
+  }
+
+  /** Stage several admissions with one membership operation per layer. */
+  prepareAppendMany(rows: readonly (readonly Cache[])[]): PreparedStateChange {
+    const next: Layout[] = [];
+    try {
+      for (let layer = 0; layer < this.caches.length; layer++) {
+        const layout = this.caches[layer]!;
+        const merged = layout.makeEmptyBatch();
+        next.push(merged);
+        merged.mergeRows([...(this.#rows ? [layout] : []), ...rows.map(row => row[layer]!)]);
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    return {
+      commit: () => {
+        const previous = this.caches.splice(0, this.caches.length, ...next.splice(0));
+        this.#rows += rows.length;
+        next.push(...previous);
+      },
+      dispose: () => disposeResources(next.splice(0)),
+    };
+  }
+
+  /** The method calls this only between committed steps. A layer failure
+   * invalidates the group; its caller discards the state. */
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.clear(); return; }
+    for (const layout of this.caches) layout.filterRows(keep);
+    this.#rows = keep.length;
+  }
+
+  /** Returns independent request state for a cache publisher or a new group. */
+  extractRow(row: number): Cache[] {
+    const state: Cache[] = [];
+    try { for (const layout of this.caches) state.push(layout.extractRow(row)); return state; }
+    catch (error) { return cleanupFailure(error, () => disposeResources(state)); }
+  }
+
+  clear(): void {
+    const next: Layout[] = [];
+    try { for (const layout of this.caches) next.push(layout.makeEmptyBatch()); }
+    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    const previous = this.caches.splice(0, this.caches.length, ...next);
+    this.#rows = 0;
+    disposeResources(previous);
+  }
+
+  dispose(): void { this.#rows = 0; disposeResources(this.caches); }
+}

package/src/state/softcap-attention.ts

@@ -0,0 +1,118 @@
+// The score-softcapped attention read (mlx-lm gemma2.py's manual attention):
+// a graph whose layers softcap their attention scores composes it at
+// construction in place of the cache's own read. The fused SDPA kernel cannot
+// express the softcap, so this read fetches the window as plain keys and values
+// and attends with the manual composition.
+//
+// It reads every KV storage through its dense fetch (`updateAndFetch`): plain
+// and rotating storage and their row layouts hand their keys and values out as
+// stored, TurboQuant storage decoded, affine storage dequantized (the delayed
+// layouts per row, before and after a row's transition), paged storage
+// gathered. Affine storage is dequantized and then read by the same manual
+// attention, so a converted layer keeps the graph's own arithmetic over the
+// values its storage holds. The read takes the storage's mask for the query
+// count before the append, appends through `updateAndFetch`, then attends
+// manually under that mask. The fetch and the mask are storage reads internal
+// to `state/` (`KvStorage`); the cache contract carries neither.
+
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { AttentionRead, Cache, Mask } from "../contracts/mlx/cache";
+import { createCausalMask } from "../kernels/attention/masks";
+import { logitSoftcap } from "../kernels/logits";
+import { disposing } from "../layers/helpers";
+import { FINFO_MIN } from "../layers/quantized-attention";
+import { disposeResources } from "../runtime/resources";
+import { kvStorage, type KvStorage } from "./kv-storage";
+
+/** The softcapped reads of one attention layer, each appending `k` and `v`
+ * (borrowed; positional encoding applied) to `cache` for the phase it names and
+ * returning the read the layer attends once and disposes. `attend(q, scale)`
+ * takes queries [B, Hq, L, D] (borrowed) and returns an owned [B, Hq, L, Dv]. */
+export interface SoftcapAttentionReads {
+  /** One new position per row. */
+  appendDecode(cache: Cache, k: MlxArray, v: MlxArray): AttentionRead;
+  /** `L` new positions per row, causal within the window. */
+  appendWindow(cache: Cache, k: MlxArray, v: MlxArray): AttentionRead;
+}
+
+/** The softcapped reads of a layer whose attention scores are capped at
+ * `softcap` (`tanh(scores / softcap) * softcap`). `window` is the layer's
+ * sliding window (null for a full-attention layer), the window its cache masks
+ * with. `arrayMasks` materializes the ordinary causal mask as a boolean array
+ * (gemma2's `create_attention_mask(..., return_array=True)`); the manual
+ * attention applies only array masks. Decode and window make the same call:
+ * the cache's own mask for the query count, which is no mask for one unpadded
+ * position. */
+export function softcapAttentionReads(softcap: number, window: number | null, arrayMasks: boolean): SoftcapAttentionReads {
+  const mask = (cache: KvStorage, L: number): Mask => {
+    const own = cache.makeMask(L, window);
+    return arrayMasks && own.mode === "causal" ? { mode: "array", arr: createCausalMask(L, cache.offset, window) } : own;
+  };
+  const read = (row: Cache, k: MlxArray, v: MlxArray): AttentionRead => {
+    const cache = kvStorage(row);
+    const held = mask(cache, k.shape[2]!);
+    let keys: MlxArray, values: MlxArray;
+    try { [keys, values] = cache.updateAndFetch(k, v); }
+    catch (error) { held.arr?.dispose(); throw error; }
+    return {
+      attend: (q, scale) => softcapAttention(q, keys, values, held, scale, softcap),
+      dispose() { disposeResources(held.arr ? [keys, values, held.arr] : [keys, values]); },
+    };
+  };
+  return Object.freeze({ appendDecode: read, appendWindow: read });
+}
+
+/** gemma2.py's manual attention: scale, the GQA 5-d reshape, q·kᵀ, the tanh
+ * softcap, a boolean mask through `where(mask, scores, finfo.min)`, the
+ * precise softmax, then ·v. `q` [B, Hq, L, D], `keys` [B, Hkv, S, D] and
+ * `values` [B, Hkv, S, Dv] are borrowed; only an array mask is applied. Returns
+ * an owned [B, Hq, L, Dv]. */
+export function softcapAttention(q: MlxArray, keys: MlxArray, values: MlxArray, mask: Mask, scale: number,
+  softcap: number): MlxArray {
+  const [B, numHeads, L, headDim] = q.shape as [number, number, number, number];
+  const numKvHeads = keys.shape[1]!;
+  const repeats = Math.floor(numHeads / numKvHeads);
+  let queries = ops.mulScalar(q, scale); // queries * self.scale
+
+  let kk = keys;
+  let vv = values;
+  let kExp: MlxArray | null = null;
+  let vExp: MlxArray | null = null;
+  if (repeats > 1) {
+    queries = disposing(queries, ops.reshape(queries, [B, numKvHeads, repeats, L, headDim]));
+    kExp = ops.expandDims(keys, 2);
+    vExp = ops.expandDims(values, 2);
+    kk = kExp;
+    vv = vExp;
+  }
+
+  const kT = ops.transposeAxes(kk, repeats > 1 ? [0, 1, 2, 4, 3] : [0, 1, 3, 2]);
+  let scores = ops.matmul(queries, kT);
+  kT.dispose();
+  queries.dispose();
+
+  // tanh(scores / cap) * cap — same composition as logitSoftcap.
+  scores = disposing(scores, logitSoftcap(scores, softcap));
+
+  if (mask.mode === "array" && mask.arr) {
+    // bool mask: where(mask, scores, finfo(scores.dtype).min)
+    const minVal = FINFO_MIN[scores.dtype];
+    if (minVal === undefined) throw new Error(`no finfo.min for dtype ${scores.dtypeName}`);
+    const minArr = ops.scalarLike(minVal, scores);
+    // GQA scores are [B, Hkv, repeats, L, S]. Row masks carry B,
+    // unlike the ordinary [L, S] causal mask, so keep that axis aligned.
+    using rowMask = repeats > 1 && mask.arr.shape.length === 4
+      ? ops.expandDims(mask.arr, 2) : null;
+    scores = disposing(scores, ops.where(rowMask ?? mask.arr, scores, minArr));
+    minArr.dispose();
+  }
+
+  scores = disposing(scores, ops.softmaxAxis(scores, -1, true));
+  let out = ops.matmul(scores, vv);
+  scores.dispose();
+  kExp?.dispose();
+  vExp?.dispose();
+  if (repeats > 1) out = disposing(out, ops.reshape(out, [B, numHeads, L, values.shape[3]!]));
+  return out;
+}

package/src/state/speculative/glm52-mtp-state.ts

@@ -0,0 +1,38 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { MLACache } from "../glm52-cache";
+import { cloneAttachments } from "../checkpoint";
+import { disposeResources } from "../../runtime/resources";
+import { DRAFT_CHECKPOINT_SCHEMA, type DraftRowCheckpoint } from "../../contracts/mlx/draft-checkpoint";
+
+export function captureGlm52MtpState(state: Glm52MtpRowState): DraftRowCheckpoint {
+  const held: MlxArray[] = [];
+  try {
+    if (state.cache.offset) {
+      const cache = state.cache.fetch(); held.push(cache.latent, cache.rope);
+    }
+    held.push(state.hidden.slice([0, 0, 0], [...state.hidden.shape]));
+    return { processedTokens: state.processedTokens, attachment: {
+      schema: DRAFT_CHECKPOINT_SCHEMA.glm52Mtp, metadata: { processedTokens: state.processedTokens, draftOffset: state.cache.offset },
+      tensors: held.splice(0),
+    } };
+  } finally { disposeResources(held); }
+}
+
+export function restoreGlm52MtpState(checkpoint: DraftRowCheckpoint, makeCache: () => MLACache): Glm52MtpRowState {
+  const { attachment, processedTokens } = checkpoint;
+  if (attachment.schema !== DRAFT_CHECKPOINT_SCHEMA.glm52Mtp || attachment.metadata.processedTokens !== processedTokens)
+    throw new Error("invalid native GLM MTP checkpoint alignment");
+  const held = cloneAttachments([attachment])[0]!.tensors;
+  let cache: MLACache | null = makeCache();
+  try {
+    const offset = attachment.metadata.draftOffset as number;
+    if (offset) { cache.restoreCompressedState(held[0]!, held[1]!, null, offset); held.splice(0, 2); }
+    const state = { cache, hidden: held.pop()!, processedTokens }; cache = null; return state;
+  } finally { disposeResources([...held, ...(cache ? [cache] : [])]); }
+}
+
+export interface Glm52MtpRowState {
+  readonly cache: MLACache;
+  readonly hidden: MlxArray;
+  readonly processedTokens: number;
+}

package/src/state/speculative/qwen-mtp-state.ts

@@ -0,0 +1,34 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { KVCache } from "../kv";
+import { cloneAttachments } from "../checkpoint";
+import { disposeResources } from "../../runtime/resources";
+import { DRAFT_CHECKPOINT_SCHEMA, type DraftRowCheckpoint } from "../../contracts/mlx/draft-checkpoint";
+
+/** The existing Qwen companion attachment is also the state interchange
+ * between prepared requests and a group. Storage-tier policy stays outside. */
+export function restoreQwenMtpState(checkpoint: DraftRowCheckpoint): MtpRowState {
+  const { processedTokens: tokens, attachment } = checkpoint;
+  if (attachment.schema !== DRAFT_CHECKPOINT_SCHEMA.qwenMtp || attachment.metadata.draftOffset !== tokens - 1)
+    throw new Error("invalid paired Qwen MTP checkpoint alignment");
+  const held = cloneAttachments([attachment])[0]!.tensors;
+  let cache: KVCache | null = new KVCache();
+  try {
+    if (tokens > 1) { cache.restoreState(held[0]!, held[1]!, tokens - 1); held.splice(0, 2); }
+    const state = { cache, hidden: held.pop()! }; cache = null; return state;
+  } finally { disposeResources([...held, ...(cache ? [cache] : [])]); }
+}
+
+export function captureQwenMtpState(state: MtpRowState): DraftRowCheckpoint {
+  const held: MlxArray[] = [];
+  try {
+    if (state.cache.offset > 0) held.push(...state.cache.temporalView());
+    held.push(state.hidden.slice([0, 0, 0], [...state.hidden.shape]));
+    return { processedTokens: state.cache.offset + 1,
+      attachment: { schema: DRAFT_CHECKPOINT_SCHEMA.qwenMtp, metadata: { draftOffset: state.cache.offset }, tensors: held.splice(0) } };
+  } finally { disposeResources(held); }
+}
+
+export interface MtpRowState {
+  readonly cache: KVCache;
+  readonly hidden: MlxArray;
+}

package/src/state/speculative-rotating-kv.ts

@@ -0,0 +1,113 @@
+import { rotatingSourcePosition } from "./rotating-kv-layout";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { type AttentionRead, type Cache, type CommittedAttentionCache, type Mask, type PrefillPadding, type SpeculativeRowLayout } from "../contracts/mlx/cache";
+import { RotatingKVCache } from "./rotating-kv";
+import { BatchedRotatingCache } from "./batched-rotating";
+import { rollbackRotatingRing } from "./rotating-row-transaction";
+import { BatchedRotatingState } from "./batched-rotating-state";
+import { plainRowStorage, temporalStorageView } from "./batched-row-storage";
+import { disposeResources } from "../runtime/resources";
+import { AttentionMasks, decodeByPosition } from "./attention-read";
+
+/** Ring transactions retain accepted KV columns, including rounds that cross
+ * the window boundary. The verify block already retains the complete history
+ * needed by its earliest query; no target recomputation or scheduler policy.
+ * Reads are the current ring's (`BatchedRotatingCache`); every ring this layout
+ * builds shares its `masks`. */
+export class SpeculativeRotatingKVCache implements CommittedAttentionCache, SpeculativeRowLayout {
+  #inner: BatchedRotatingCache;
+  #before?: number[];
+  constructor(
+    readonly maxSize: number,
+    /** The per-forward masks of the model this layout serves. */
+    readonly masks = new AttentionMasks(),
+  ) { this.#inner = new BatchedRotatingCache(maxSize, [], masks); }
+  signature(): string { return "kv:rotating-plain"; }
+  get offset(): number { return this.#inner.offset; }
+  get rowOffsets(): readonly number[] { return this.#inner.offsetArr; }
+  get leftPad(): readonly number[] { return this.#inner.leftPad; }
+  get batchSize(): number { return this.#inner.batchSize; }
+  get ropeOffsetArr(): MlxArray { return this.#inner.ropeOffsetArr; }
+  makeEmptyBatch(): SpeculativeRotatingKVCache { return new SpeculativeRotatingKVCache(this.maxSize, this.masks); }
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead { return this.#inner.appendDecode(k, v); }
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead { return this.#inner.appendWindow(k, v); }
+  /** A committed span, one decode per position (`decodeByPosition`). */
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead { return decodeByPosition(this, k, v); }
+  makeMask(tokens: number, window: number | null): Mask { return this.#inner.makeMask(tokens, window); }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#inner.updateAndFetch(k, v); }
+  preparePrefill(padding: PrefillPadding): void { this.#inner.preparePrefill(padding); }
+  finalizePrefill(): void { this.#inner.finalizePrefill(); }
+  state(): MlxArray[] { return this.#inner.state(); }
+  isTrimmable(): boolean { return this.#inner.isTrimmable(); }
+  trim(count: number): void { this.#inner.trim(count); }
+  bytesPerToken(): number {
+    const state = this.state();
+    return state.length ? state.reduce((n, a) => n + a.nbytes, 0) / (this.batchSize * state[0]!.shape[2]!) : 0;
+  }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * Math.min(tokens, this.maxSize); }
+  extractRow(row: number): Cache {
+    const result = new RotatingKVCache(this.maxSize);
+    if (!this.rowOffsets[row] || !this.#inner.keys || !this.#inner.values) return result;
+    const position = new BatchedRotatingState(this.maxSize, []); position.restore(this.#inner.positionSnapshot);
+    const range = { row, from: Math.max(0, position.leftPad[row]!, position.activeLength - this.maxSize),
+      to: position.activeLength, copy: true };
+    const keys = temporalStorageView(plainRowStorage, this.#inner.keys, position, range);
+    let values: MlxArray;
+    try { values = temporalStorageView(plainRowStorage, this.#inner.values, position, range); }
+    catch (error) { keys.dispose(); throw error; }
+    result.restoreState(keys, values, this.rowOffsets[row]!, keys.shape[2]!); return result;
+  }
+  captureDonorRows(): import("../contracts/mlx/cache").KvDonorRows {
+    const position = new BatchedRotatingState(this.maxSize, []); position.restore(this.#inner.positionSnapshot);
+    const from = Math.max(0, position.activeLength - this.maxSize);
+    const range = { from, to: position.activeLength };
+    const keys = temporalStorageView(plainRowStorage, this.#inner.keys!, position, range); let values: MlxArray;
+    try { values = temporalStorageView(plainRowStorage, this.#inner.values!, position, range); }
+    catch (error) { keys.dispose(); throw error; }
+    return { keys, values, offsets: [...this.rowOffsets], starts: this.leftPad.map(pad => Math.max(0, pad - from)),
+      ends: this.rowOffsets.map(() => position.activeLength - from) };
+  }
+
+  filterRows(keep: readonly number[]): void { this.#inner.filterRows(keep); }
+  mergeRows(sources: readonly Cache[]): void {
+    const only = sources.length === 1 ? sources[0] : undefined;
+    if (only instanceof SpeculativeRotatingKVCache) {
+      const keys = only.#inner.keys ? ops.copyOf(only.#inner.keys) : null;
+      let values: MlxArray | null;
+      try { values = only.#inner.values ? ops.copyOf(only.#inner.values) : null; }
+      catch (error) { keys?.dispose(); throw error; }
+      const next = BatchedRotatingCache.adoptPhysical(keys, values, only.#inner.positionSnapshot, this.masks);
+      this.#inner.dispose(); this.#inner = next; return;
+    }
+    const held: Cache[] = [], planes: { keys: MlxArray; values: MlxArray }[] = [], offsets: number[] = [];
+    try {
+      const rows = sources.flatMap(source => source instanceof SpeculativeRotatingKVCache
+        ? Array.from({ length: source.batchSize }, (_, row) => { const cache = source.extractRow(row); held.push(cache); return cache; })
+        : [source]) as RotatingKVCache[];
+      const prototype = rows.find(row => row.keys !== null);
+      if (!prototype) {
+        this.#inner.dispose(); this.#inner = new BatchedRotatingCache(this.maxSize, rows.map(() => 0), this.masks); return;
+      }
+      for (const cache of rows) {
+        const state = rotatingSourcePosition(cache);
+        const range = { from: Math.max(0, state.activeLength - this.maxSize), to: state.activeLength };
+        const plane = (field: "keys" | "values") => cache[field]
+          ? temporalStorageView(plainRowStorage, cache[field]!, state, range)
+          : ops.zeros([1, prototype[field]!.shape[1]!, 0, prototype[field]!.shape[3]!], prototype[field]!.dtype);
+        const keys = plane("keys"); let values: MlxArray;
+        try { values = plane("values"); } catch (error) { keys.dispose(); throw error; }
+        planes.push({ keys, values }); offsets.push(cache.offset);
+      }
+      const next = BatchedRotatingCache.merge(planes, offsets, this.maxSize, this.masks);
+      this.#inner.dispose(); this.#inner = next;
+    } finally { disposeResources([...held, ...planes.flatMap(row => [row.keys, row.values])]); }
+  }
+  specRoundBegin(): void { this.#before = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#before = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    this.#inner = rollbackRotatingRing(this.#inner, this.#before!, keep);
+    this.#before = undefined;
+  }
+  dispose(): void { this.#before = undefined; this.#inner.dispose(); }
+}