@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/scoring/full-sequence.ts

@@ -0,0 +1,103 @@
+// Training forward pass: full-sequence logits with no KV reuse.
+//
+// Every forward runs the graph's scoring forward (sequenceHidden) over fresh
+// stateless training caches (state/training-cache.ts), one per cache the graph
+// composes: a TrainingCache with that layer's sliding window for attention, a
+// TrainingSSMCache for recurrent state. They share one mask memo and attend
+// with the training attention the run selected; the caches own their reads, so
+// nothing here builds a mask or reads keys and values.
+//
+// B=1 (no padding): offset-0 causal (windowed) reads, safe to recompute under
+// gradient checkpointing — bit-identical to the single-example trainer.
+//
+// B>1 (padded batch): every row in the batch is padded to a common length L,
+// but the rows have different *true* lengths. A plain causal mask would let a
+// padded query attend padded keys, so the caches take the rows' lengths and
+// read under the padded-batch mask: causal (within the window) AND key within
+// its row's true length, built at offset 0 like the unpadded read.
+
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { type Cache } from "../contracts/mlx/cache";
+import { AttentionMasks } from "../state/attention-read";
+import { selectedTrainingAttention, TrainingCache, TrainingSSMCache, trainingStandIn } from "../state/training-cache";
+import type { MlxSequenceForward, MlxSequenceRange } from "../contracts/mlx/graph";
+import type { RuntimeModel } from "../models/factory";
+
+export { TrainingCache, TrainingSSMCache };
+/** The padded-batch training mask (`paddedTrainingMask`): [B, 1, L, L] bool,
+ *  causal (within `windowSize`) AND key within its row's valid length. Caller
+ *  owns the result. */
+export { paddedTrainingMask as buildBatchedPadMask } from "../state/training-cache";
+
+/** Run a full-sequence forward for training.
+ *  @param ids int32 array [B, L].
+ *  @returns logits [B, L, V] (caller owns; dispose when done).
+ *
+ *  When `validLengths` is omitted (or every row is full-length L) the batch
+ *  has no padding and the caches read causally — the B=1 result is
+ *  bit-identical to the original trainer. When some row is shorter than L,
+ *  the caches read under the padding-aware batched mask. */
+export function trainForwardHidden(
+  model: RuntimeModel,
+  ids: MlxArray,
+  validLengths?: number[],
+): MlxArray {
+  const [B, L] = ids.shape as [number, number];
+  const width = (B * L);
+  const rowLengths = validLengths !== undefined && validLengths.some((v) => v < L) ? validLengths : null;
+  // Stateless offset-0 caches (one per cache the graph composes), so the
+  // forward is pure and safe to recompute under gradient checkpointing. Each
+  // keeps its layer's sliding window; recurrent layers take TrainingSSMCache,
+  // whose causal recurrence never carries a right-padded position into a real
+  // one, so it needs no mask.
+  const probe = model.makeCache();
+  const masks = new AttentionMasks(), attention = selectedTrainingAttention();
+  const cache: Cache[] = probe.map((c) =>
+    trainingStandIn(c, window => new TrainingCache(window, masks, attention, rowLengths)));
+  for (const c of probe) c.dispose();
+  try {
+    return sequenceHiddenOperation(model, ids, cache, width); // [B, L, hidden]
+  } finally {
+    for (const c of cache) c.dispose();
+  }
+}
+
+/** Full-sequence forward returning logits [B, L, V]. Convenience over
+ *  trainForwardHidden + logitsFromHidden. The SFT B=1 path deliberately skips
+ *  this and applies the LM head only at response positions (responseOnlyCe)
+ *  to avoid materializing the [B, L, vocab] logits — the dominant training
+ *  memory term at long context with a 262k vocab. */
+export function trainForward(
+  model: RuntimeModel,
+  ids: MlxArray,
+  validLengths?: number[],
+): MlxArray {
+  const h = trainForwardHidden(model, ids, validLengths);
+  const logits = sequenceLogitsOperation(model,h);
+  h.dispose();
+  return logits;
+}
+
+
+/** Select the full-sequence operation from its exact token rows. */
+type SequenceOperations = { readonly sequence?: readonly MlxSequenceRange[]; readonly sequenceHidden?: MlxSequenceForward };
+export function sequenceHiddenOperation(model: SequenceOperations | object,
+  ids: MlxArray, cache: Cache[], rows = ids.shape[0]! * ids.shape[1]!): MlxArray {
+  const ports = model as SequenceOperations;
+  if (ports.sequence) return ports.sequence.find(range => rows >= range.from && rows <= range.to)!.operation(ids, cache);
+  if (!ports.sequenceHidden) throw new TypeError("the graph declares no full-sequence operation");
+  return ports.sequenceHidden(ids, cache, rows);
+}
+
+/** Full-sequence heads may cover more rows than the serving scheduler's table. */
+export function sequenceLogitsOperation(model: object, hidden: MlxArray): MlxArray {
+  const ports=model as {readonly sequenceLogits?: readonly import("../contracts/mlx/graph").MlxHeadRange[]; logitsFromHidden?(hidden:MlxArray):MlxArray};
+  const rows=hidden.shape.slice(0,-1).reduce((a,b)=>a*b,1);
+  if (ports.sequenceLogits) {
+    const range=ports.sequenceLogits.find(range=>rows>=range.from && rows<=range.to);
+    if (!range) throw new Error(`graph has no sequence head for ${rows} rows`);
+    return range.operation(hidden);
+  }
+  if (ports.logitsFromHidden) return ports.logitsFromHidden(hidden);
+  throw new TypeError("graph declares no full-sequence head");
+}

package/src/scoring/index.ts

@@ -0,0 +1,9 @@
+export * from "../kernels/logits";
+export * from "../kernels/softcap";
+export * from "./full-sequence";
+export { trainForward as forwardSequence, trainForwardHidden as forwardSequenceHidden } from "./full-sequence";
+export * from "./perplexity";
+export * from "./kl";
+export * from "./decisions";
+export * from "./response-logps";
+

package/src/scoring/kl.ts

@@ -0,0 +1,29 @@
+import { MlxArray } from '@mlx-bun/mlx/array';
+import { Dtype } from '@mlx-bun/mlx/ffi';
+import * as ops from '@mlx-bun/mlx/ops';
+
+
+/** KL(p ‖ q) per token for logits [1, T, V]; returns a length-T array. */
+export function klPerToken(pLogits: MlxArray, qLogits: MlxArray): Float32Array {
+  const f32 = (a: MlxArray): { arr: MlxArray; owned: boolean } =>
+    a.dtype === Dtype.float32 ? { arr: a, owned: false } : { arr: a.astype(Dtype.float32), owned: true };
+  const pp = f32(pLogits);
+  const qq = f32(qLogits);
+  const p = pp.arr;
+  const q = qq.arr;
+
+  const lseP = ops.logsumexpAxis(p, -1, true);   // [1,T,1]
+  const lseQ = ops.logsumexpAxis(q, -1, true);
+  const logP = ops.sub(p, lseP);                  // [1,T,V]
+  const logQ = ops.sub(q, lseQ);
+  const pProbs = ops.softmaxAxis(p, -1, true);    // precise softmax
+  const diff = ops.sub(logP, logQ);
+  const prod = ops.mul(pProbs, diff);
+  const kl = ops.sumAxis(prod, -1, false);        // [1,T]
+  const out = kl.toFloat32();
+
+  for (const a of [lseP, lseQ, logP, logQ, pProbs, diff, prod, kl]) a.dispose();
+  if (pp.owned) p.dispose();
+  if (qq.owned) q.dispose();
+  return out;
+}

package/src/scoring/perplexity.ts

@@ -0,0 +1,165 @@
+// Perplexity over a local text/JSONL dataset — mlx_lm.perplexity methodology.
+//
+// The reference (mlx_lm/perplexity.py, mlx-lm 0.31.3) does:
+//   1. tokenize dataset samples, visit them in a seeded random permutation,
+//   2. concatenate tokens into one stream and cut it into NON-OVERLAPPING
+//      rows of `sequence_length` (no sliding-window stride, no BOS re-seeding
+//      per row), keeping at most `num_samples` rows,
+//   3. per batch: logits = model(batch[:, :-1]) in f32, per-token
+//      cross-entropy vs batch[:, 1:] — EVERY position counts,
+//   4. ppl = exp(mean loss); SE via the delta method:
+//      se_ppl = ppl · std(losses, ddof=1) / √N.
+//
+// We reproduce 2-4 exactly. Step 1 differs only in the data source: the
+// reference pulls a Hugging Face dataset (network); we take a LOCAL .txt or
+// .jsonl file — never a download. The permutation RNG is a seeded
+// deterministic shuffle (not NumPy's), so row selection is reproducible per
+// seed but not bit-matched to the Python tool's sampling.
+//
+// The forward reuses the existing full-sequence machinery
+// (packages/inference/src/scoring/full-sequence.ts trainForward →
+// the graph's sequenceHidden + logitsFromHidden).
+
+import { clearCache, Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { trainForward } from "./full-sequence";
+import type { RuntimeModel } from "../models/factory";
+
+export interface PplDataOptions {
+  sequenceLength: number;
+  /** Max rows to keep (-1 / 0 = all available). */
+  numSamples: number;
+  /** Seed for the sample-order shuffle (reference default: 123). */
+  seed: number;
+}
+
+export interface PplResult {
+  ppl: number;
+  /** Delta-method standard error of the perplexity. */
+  standardError: number;
+  meanLoss: number;
+  tokens: number;
+  rows: number;
+}
+
+/** Deterministic PRNG (mulberry32) for the seeded sample shuffle. */
+function mulberry32(seed: number): () => number {
+  let s = seed >>> 0;
+  return () => {
+    s = (s + 0x6d2b79f5) >>> 0;
+    let t = s;
+    t = Math.imul(t ^ (t >>> 15), t | 1);
+    t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
+    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
+  };
+}
+
+/** Split a raw dataset file into text samples: `.jsonl` → one sample per
+ *  line from the `text` field; anything else → the whole file is one sample. */
+export function parseSamples(raw: string, path: string): string[] {
+  if (path.endsWith(".jsonl") || path.endsWith(".ndjson")) {
+    const samples: string[] = [];
+    for (const [i, line] of raw.split("\n").entries()) {
+      if (!line.trim()) continue;
+      let row: unknown;
+      try {
+        row = JSON.parse(line);
+      } catch {
+        throw new Error(`${path}:${i + 1}: not valid JSON`);
+      }
+      const text = (row as { text?: unknown }).text;
+      if (typeof text !== "string")
+        throw new Error(`${path}:${i + 1}: expected a {"text": "..."} row`);
+      samples.push(text);
+    }
+    return samples;
+  }
+  return raw.length > 0 ? [raw] : [];
+}
+
+/** Tokenize samples (seeded shuffle), concatenate, and pack into
+ *  non-overlapping rows of `sequenceLength` tokens (reference load_data). */
+export function packRows(
+  sampleTokens: number[][],
+  opts: PplDataOptions,
+): Int32Array[] {
+  const { sequenceLength: L, numSamples, seed } = opts;
+  const order = sampleTokens.map((_, i) => i);
+  const rnd = mulberry32(seed);
+  for (let i = order.length - 1; i > 0; i--) {
+    const j = Math.floor(rnd() * (i + 1));
+    [order[i], order[j]] = [order[j]!, order[i]!];
+  }
+
+  const wanted = numSamples > 0 ? L * numSamples : Infinity;
+  const stream: number[] = [];
+  for (const idx of order) {
+    if (stream.length >= wanted) break;
+    stream.push(...sampleTokens[idx]!);
+  }
+
+  const nRows = Math.min(
+    Math.floor(stream.length / L),
+    numSamples > 0 ? numSamples : Infinity,
+  );
+  const rows: Int32Array[] = [];
+  for (let r = 0; r < nRows; r++)
+    rows.push(Int32Array.from(stream.slice(r * L, (r + 1) * L)));
+  return rows;
+}
+
+/** Evaluate perplexity over packed rows. Per batch (all rows are the same
+ *  length, so the stock unpadded forward applies): f32 logits over
+ *  batch[:, :-1], per-token CE vs batch[:, 1:], every position counted. */
+export function evalPpl(
+  model: RuntimeModel,
+  rows: Int32Array[],
+  batchSize: number,
+  onBatch?: (done: number, total: number) => void,
+): PplResult {
+  const L = rows[0]!.length;
+  const allLosses: number[] = [];
+  const numBatches = Math.ceil(rows.length / batchSize);
+
+  for (let s = 0, bi = 0; s < rows.length; s += batchSize, bi++) {
+    const batch = rows.slice(s, s + batchSize);
+    const B = batch.length;
+    const inputs: number[] = [];
+    const targets: number[] = [];
+    for (const row of batch) {
+      for (let i = 0; i < L - 1; i++) inputs.push(row[i]!);
+      for (let i = 1; i < L; i++) targets.push(row[i]!);
+    }
+    const ids = ops.fromInt32(inputs, [B, L - 1]);
+    const tgt = ops.fromInt32(targets, [B, L - 1, 1]);
+    const logits = trainForward(model, ids);
+    const f32 = logits.astype(Dtype.float32);
+    logits.dispose();
+    // CE[t] = logsumexp(logits[t]) − logits[t, target]  (f32, like nn.losses)
+    const lse = ops.logsumexpAxis(f32, 2, false); // [B, L-1]
+    const picked = ops.takeAlongAxis(f32, tgt, 2); // [B, L-1, 1]
+    f32.dispose();
+    const picked2d = ops.reshape(picked, [B, L - 1]);
+    picked.dispose();
+    const losses = ops.sub(lse, picked2d);
+    lse.dispose();
+    picked2d.dispose();
+    const vals = losses.toFloat32();
+    losses.dispose();
+    ids.dispose();
+    tgt.dispose();
+    for (const v of vals) allLosses.push(v);
+    clearCache();
+    onBatch?.(bi + 1, numBatches);
+  }
+
+  const n = allLosses.length;
+  const mean = allLosses.reduce((a, b) => a + b, 0) / n;
+  const ppl = Math.exp(mean);
+  // Sample std (ddof=1) → delta-method SE of the perplexity.
+  const varSum = allLosses.reduce((a, b) => a + (b - mean) ** 2, 0);
+  const std = n > 1 ? Math.sqrt(varSum / (n - 1)) : 0;
+  const standardError = ppl * (std / Math.sqrt(n));
+
+  return { ppl, standardError, meanLoss: mean, tokens: n, rows: rows.length };
+}

package/src/scoring/response-logps.ts

@@ -0,0 +1,65 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { sequenceLogitsOperation } from "./full-sequence";
+
+/** The existing DPO/ORPO head geometry, shared with inference. Two- and
+ * three-row response spans pad to four before the head and slice back. Borrows
+ * hidden and model, returns caller-owned logits; preserves training numerics. */
+export function logitsFromHiddenPadM(
+  model: object,
+  hSpan: MlxArray, // [1, M, hidden]
+): MlxArray {
+  const M = hSpan.shape[1]!;
+  if (M !== 2 && M !== 3) return sequenceLogitsOperation(model, hSpan);
+  const [B, , H] = hSpan.shape as [number, number, number];
+  const pad = ops.zeros([B, 4 - M, H], hSpan.dtype);
+  const hPad = ops.concatAxis([hSpan, pad], 1);
+  pad.dispose();
+  const lgPad = sequenceLogitsOperation(model, hPad); // [1, 4, V]
+  hPad.dispose();
+  const start = MlxArray.fromInt32(new Int32Array([0]), [1]);
+  const lg = ops.sliceDynamic(lgPad, start, [1], [B, M, lgPad.shape[2]!]);
+  start.dispose();
+  lgPad.dispose();
+  return lg;
+}
+
+/** Per-row summed response log likelihood, extracted unchanged from DPO.
+ * logits [B,L-1,V] predict ids[:,1:]; mask is indexed in the original [B,L]
+ * token coordinates. Prompt and padding masks are zero, answer masks one.
+ * Borrows every input and returns an owned differentiable f32 tensor [B].
+ * Subtraction precedes the f32 cast, retaining the existing DPO operation order.
+ */
+export function sequenceLogProbabilities(logits: MlxArray, ids: number[][], mask: number[][]): MlxArray {
+  const B = ids.length;
+  const L = ids[0]!.length;
+  const T = L - 1;
+  const V = logits.shape[2]!;
+  const N = B * T;
+
+  const targetsHost = new Int32Array(N);
+  const maskHost = new Float32Array(N);
+  for (let b = 0; b < B; b++)
+    for (let t = 0; t < T; t++) {
+      const flat = b * T + t;
+      targetsHost[flat] = ids[b]![t + 1]!;
+      maskHost[flat] = mask[b]![t + 1]!; // resp_mask[:, 1:]
+    }
+  const logits2d = ops.reshape(logits, [N, V]);
+  const targets = MlxArray.fromInt32(targetsHost, [N, 1]);
+  const m = MlxArray.fromFloat32(maskHost, [N]);
+
+  const lse = ops.logsumexpAxis(logits2d, -1, false); // [N]
+  const gathered = ops.takeAlongAxis(logits2d, targets, -1); // [N,1]
+  const picked = ops.reshape(gathered, [N]); // [N]
+  const logp = ops.sub(picked, lse); // log_softmax at the label
+  const logpF = logp.dtype === Dtype.float32 ? logp : logp.astype(Dtype.float32);
+  const masked = ops.mul(logpF, m); // [N]
+  const masked2d = ops.reshape(masked, [B, T]); // [B, T]
+  const perRow = ops.sumAxis(masked2d, 1, false); // [B]
+
+  for (const a of [logits2d, targets, m, lse, gathered, picked, logp, masked, masked2d]) a.dispose();
+  if (logpF !== logp) logpF.dispose();
+  return perRow;
+}

package/src/state/affine-attention.ts

@@ -0,0 +1,180 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { AttentionRead, Mask } from "../contracts/mlx/cache";
+import type { AffineKernelChoice } from "../contracts/portable/composition";
+import { quantizedAppendAttention } from "../layers/quantized-append-attention";
+import { fusedSdpaConfigSupported, quantizedSdpaTiled, quantizedSdpaUnfused } from "../layers/quantized-attention";
+import type { AttentionMasks, MaskLease } from "./attention-read";
+import { disposeTriple } from "./quantized-tensor";
+
+/** One attention kernel over affine-quantized keys and values: queries
+ * [B, Hq, L, D] to an owned output [B, Hq, L, Dv]. Borrows every input. */
+export type AffineKernel = (q: MlxArray, keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
+  scale: number, mask: Mask) => MlxArray;
+
+/** A kernel set named by the composition it serves (`composedAffineKernels`,
+ * `unfusedAffineKernels`, `tiledCausalAffineKernels`, `tiledAffineKernels`),
+ * which a conversion builds for the converted cache's bits, group size and
+ * activation dtype. */
+export type AffineKernelSet = (bits: number, groupSize: number, dtype: Dtype) => AffineKernels;
+
+/** The kernel an affine cache attends with for each read. The cache takes it
+ * when it is built and never chooses again: `unfusedAffineKernels` for a
+ * uniform `--kv-quant N`, `tiledCausalAffineKernels` for `--kv-quant config`,
+ * `tiledAffineKernels` for the opt-in always-tiled composition, or a
+ * device-specific cache's own. */
+export interface AffineKernels {
+  /** The decode read: one query per row. */
+  readonly decode: AffineKernel;
+  /** The window read (`L` causal queries per row) while the cache's own mask
+   * is plain causal. */
+  readonly window: AffineKernel;
+  /** The window read once the cache's own state makes its mask an array:
+   * padded or unequal rows in a batched cache, the sliding window bound in a
+   * rotating cache. The cache picks it on the branch that picks that mask,
+   * from state it owns, never from anything the caller passes. */
+  readonly maskedWindow: AffineKernel;
+  /** The kernels a row layout reads with unless it holds exactly one row at
+   * its own positions (several rows, or one behind left padding); absent,
+   * these kernels at every row count. Only a set whose kernels take one
+   * sequence has it (the M4 Pro composition's KV4 set, `Kv4Head256Cache`). The
+   * layout picks it from its own row layout, as it picks `maskedWindow` from
+   * its own mask, never from the queries it is handed. */
+  readonly multiRow?: AffineKernels;
+}
+
+/** The kernels a row layout reads with: `kernels` while it holds exactly one
+ * row at its own positions (`singleUnpaddedRow`), their `multiRow` set (when
+ * they have one) otherwise. */
+export function rowLayoutKernels(kernels: AffineKernels, singleUnpaddedRow: boolean): AffineKernels {
+  return singleUnpaddedRow ? kernels : kernels.multiRow ?? kernels;
+}
+
+/** The kernel set a composition's affine caches attend with
+ * (`KvStorageScheme.kernels`): `unfused`, mlx-lm's port for every read;
+ * `tiled-causal`, OptiQ's wrapper rule where the tiled port takes the bits,
+ * group size and activation dtype and the unfused port where it does not, as
+ * the served path computed each configuration. Resolved for each cache when its
+ * quantized storage is built. */
+export function composedAffineKernels(choice: AffineKernelChoice): AffineKernelSet {
+  return choice === "unfused" ? unfusedAffineKernels
+    : (bits, groupSize, dtype) => fusedSdpaConfigSupported(groupSize, bits, dtype)
+      ? tiledCausalAffineKernels(bits, groupSize, dtype) : unfusedAffineKernels(bits, groupSize, dtype);
+}
+
+/** The read the row layouts' and donors' attention views make over fetched
+ * affine keys and values (`captureKvAttention`, `quantizedDonorAttention`):
+ * one query reads with `kernels.decode`; a window under a plain causal mask, or
+ * a causal-equivalent bool matrix, with `kernels.window`; under any other array
+ * mask with `kernels.maskedWindow`. The views choose by query count because the
+ * row layouts' attention views and the donor reads receive their queries
+ * without a phase; the caches' named reads choose by phase. */
+export function viewAffineKernel(kernels: AffineKernels, queries: number, mask: Mask): AffineKernel {
+  if (queries === 1) return kernels.decode;
+  const causal = mask.mode === "causal" || mask.mode === "" ||
+    (mask.mode === "array" && mask.causalEquivalent === true && mask.arr !== null && mask.arr.ndim === 2 && mask.arr.dtype === Dtype.bool);
+  return causal ? kernels.window : kernels.maskedWindow;
+}
+
+/** The stock unfused port (`quantizedSdpaUnfused`). */
+export function unfusedKernel(groupSize: number, bits: number): AffineKernel {
+  return (q, keys, values, scale, mask) => quantizedSdpaUnfused(q, keys, values, scale, mask, groupSize, bits);
+}
+
+/** The N-tiled fused port (`quantizedSdpaTiled`). It folds GQA into 5-D
+ * scores, so a per-row mask [B, 1, L, N] takes the unit KV-head axis there,
+ * as the unfused port inserts it for itself. */
+export function tiledKernel(groupSize: number, bits: number): AffineKernel {
+  return (q, keys, values, scale, mask) => {
+    if (mask.arr?.ndim === 4 && q.shape[1] !== keys.packed.shape[1]) {
+      using arr = ops.expandDims(mask.arr, 1);
+      return quantizedSdpaTiled(q, keys, values, scale, { mode: "array", arr }, groupSize, bits);
+    }
+    return quantizedSdpaTiled(q, keys, values, scale, mask, groupSize, bits);
+  };
+}
+
+/** Every read unfused: mlx-lm's quantized SDPA port. The composition of a
+ * uniform `--kv-quant N` (KV scheme `affine`, `kernels: "unfused"`), which
+ * keeps mlx-lm parity, and of `--fused-sdpa off`. The unfused port takes any
+ * bits, group size and activation dtype, so the dtype only names the
+ * configuration. */
+export function unfusedAffineKernels(bits: number, groupSize: number, _dtype: Dtype): AffineKernels {
+  const unfused = unfusedKernel(groupSize, bits);
+  return { decode: unfused, window: unfused, maskedWindow: unfused };
+}
+
+/** OptiQ's wrapper rule: decode unfused (what one query gets in OptiQ's
+ * wrapper); a window through the N-tiled fused port while the cache's own mask
+ * is plain causal, and through the unfused port once the cache's own state
+ * makes that mask an array (padded or unequal rows in a batched cache, the
+ * sliding window bound in a rotating cache). The choice keys only on state the
+ * cache owns, like the ring wrap or the delayed-quantization transition, never
+ * on anything the caller passes. The composition of `--kv-quant config` (KV
+ * scheme `affine-layers`, `kernels: "tiled-causal"`), which keeps OptiQ parity,
+ * and of `--fused-sdpa on` (`composedAffineKernels`). Throws at construction
+ * when the tiled port does not take these bits, group size and activation
+ * `dtype`. */
+export function tiledCausalAffineKernels(bits: number, groupSize: number, dtype: Dtype): AffineKernels {
+  supportsTiling("tiledCausalAffineKernels", bits, groupSize, dtype);
+  const unfused = unfusedKernel(groupSize, bits);
+  return { decode: unfused, window: tiledKernel(groupSize, bits), maskedWindow: unfused };
+}
+
+/** Decode unfused, every window through the N-tiled fused port, masked ones
+ * included. The always-tiled composition: no oracle computes it, so it is
+ * opt-in and gated by KL. Throws at construction when the tiled port does not
+ * take these bits, group size and activation `dtype`
+ * (`fusedSdpaConfigSupported`); it never falls back to the unfused port. */
+export function tiledAffineKernels(bits: number, groupSize: number, dtype: Dtype): AffineKernels {
+  supportsTiling("tiledAffineKernels", bits, groupSize, dtype);
+  const tiled = tiledKernel(groupSize, bits);
+  return { decode: unfusedKernel(groupSize, bits), window: tiled, maskedWindow: tiled };
+}
+
+function supportsTiling(name: string, bits: number, groupSize: number, dtype: Dtype): void {
+  if (!fusedSdpaConfigSupported(groupSize, bits, dtype))
+    throw new Error(`${name}: the tiled port takes 4- or 8-bit, group 32, 64 or 128, bf16 or f16 ` +
+      `activations; got ${bits}-bit, group ${groupSize}, dtype ${dtype}`);
+}
+
+/** A read over fetched affine keys and values: `kernel` under the leased
+ * mask. Takes ownership of `keys`, `values` and `mask`. */
+export function affineRead(keys: ops.QuantizedTensor, values: ops.QuantizedTensor, mask: MaskLease,
+  kernel: AffineKernel): AttentionRead {
+  return {
+    attend: (q, scale) => kernel(q, keys, values, scale, mask.mask),
+    dispose() { disposeTriple(keys); disposeTriple(values); mask.release(); },
+  };
+}
+
+/** The committed read over fetched affine keys and values: each position of
+ * the span attends its own causal prefix through the one-query unfused port
+ * (`quantizedAppendAttention`, the graphs' independent-rows path). One row.
+ * Takes ownership of `keys` and `values`. */
+export function committedAffineRead(keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
+  groupSize: number, bits: number): AttentionRead {
+  return {
+    attend: (q, scale) => quantizedAppendAttention(q, keys, values, scale, groupSize, bits),
+    dispose() { disposeTriple(keys); disposeTriple(values); },
+  };
+}
+
+/** The bottom-right-aligned causal matrix, bool [L, N] (query `i` sees keys
+ * `0..N-L+i`), that both affine kernels build for themselves from the
+ * "causal" mask, here built once per forward through `masks`. */
+export function affineCausalLease(masks: AttentionMasks, L: number, N: number): MaskLease {
+  return masks.lease("affine-causal", `${L}|${N}`, () => ({ mode: "array", arr: bottomRightCausal(L, N) }));
+}
+
+/** The kernels' "causal" branch, verbatim. */
+function bottomRightCausal(L: number, N: number): MlxArray {
+  const qIdx = ops.arange(N - L, N, 1, Dtype.int32);
+  const kIdx = ops.arange(0, N, 1, Dtype.int32);
+  const qCol = ops.reshape(qIdx, [L, 1]);
+  const kRow = ops.reshape(kIdx, [1, N]);
+  const mask = ops.greaterEqual(qCol, kRow);
+  for (const a of [qIdx, kIdx, qCol, kRow]) a.dispose();
+  return mask;
+}

package/src/state/attachment.ts

@@ -0,0 +1,68 @@
+import type { Cache } from "../contracts/mlx/cache";
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { legacyCacheCodecs } from "./persistence";
+import { type CacheCodecProvider, type CacheHeaderEntry, type SnapshotContext } from "./persistence-types";
+import { disposeResources } from "../runtime/resources";
+import type { CheckpointAttachment } from "../contracts/mlx/checkpoint";
+type TensorSlot = SnapshotContext["slots"][number];
+
+/** Reuse the backend cache codecs for method-owned companion state. Tensor
+ * slots address attachment indices here; the persistence tier supplies the
+ * outer file offsets and checksums. No host serialization of tensors occurs. */
+export function captureCacheAttachment(schema: string, caches: readonly Cache[],
+  metadata: CheckpointAttachment["metadata"], codecs: CacheCodecProvider = legacyCacheCodecs): CheckpointAttachment {
+  const tensors: MlxArray[] = [], headers: Array<CacheHeaderEntry | null> = [];
+  try {
+    for (const cache of caches) {
+      if (!cache.offset) { headers.push(null); continue; }
+      const slots: TensorSlot[] = [];
+      const liveSlice = (array: MlxArray, length: number, axis = 2) => {
+        const end = [...array.shape]; end[axis] = length;
+        return array.slice(array.shape.map(() => 0), end);
+      };
+      const push = (array: MlxArray, owned: boolean) => {
+        const value = owned ? array : ops.copyOf(array);
+        const off = tensors.length; tensors.push(value);
+        slots.push({ off, bytes: value.nbytes, shape: [...value.shape], dtype: value.dtype, hash: "" });
+      };
+      const context: SnapshotContext = { slots, push, liveSlice,
+        liveMlaSlice: (array, length) => liveSlice(array, length, 1),
+        pushTriple(triple, length) {
+          for (const array of [triple.packed, triple.scales, triple.biases])
+            push(length === null ? array : liveSlice(array, length), length !== null);
+        },
+      };
+      const stored = cache.composed ? cache.composed.storage() : cache;
+      const entry = codecs.forCache(stored).snapshot(stored, context);
+      if (stored.minimumReusableOffset !== undefined) entry.minimumReusableOffset = stored.minimumReusableOffset;
+      headers.push(entry);
+    }
+    return { schema, metadata: { ...metadata, cacheCodecs: codecs.id, cacheHeaders: JSON.stringify(headers) }, tensors };
+  } catch (error) { disposeResources(tensors); throw error; }
+}
+
+/** Returned caches own immutable aliases. The caller keeps an external SSD
+ * lease until its normal state materialization boundary has completed. */
+export function restoreCacheAttachment(attachment: CheckpointAttachment, makeCache: () => Cache[],
+  codecs: CacheCodecProvider = legacyCacheCodecs): Cache[] {
+  if (attachment.metadata.cacheCodecs !== codecs.id) throw new Error("companion cache codec mismatch");
+  const headers = JSON.parse(attachment.metadata.cacheHeaders as string) as Array<CacheHeaderEntry | null>;
+  const caches = makeCache(), pending: MlxArray[] = [];
+  const arr = (slot: TensorSlot) => {
+    const value = ops.copyOf(attachment.tensors[slot.off]!); pending.push(value); return value;
+  };
+  try {
+    for (const [index, entry] of headers.entries()) {
+      if (!entry) continue;
+      const cache = codecs.forHeader(entry).load(entry, { path: attachment.schema, prototype: caches[index]!, arr, grownArr: arr,
+        triple: (slots, at) => ({ packed: arr(slots[at]!), scales: arr(slots[at + 1]!), biases: arr(slots[at + 2]!) }),
+      });
+      pending.length = 0;
+      cache.minimumReusableOffset = entry.minimumReusableOffset ?? 0;
+      const previous = caches[index]; caches[index] = cache;
+      previous?.dispose();
+    }
+    return caches;
+  } catch (error) { disposeResources([...pending, ...caches]); throw error; }
+}