@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/batched-mask.ts

@@ -0,0 +1,445 @@
+// Padding-aware attention masks for batched (B>1) serving — decode side.
+//
+// Companion to packages/inference/src/scoring/full-sequence.ts's buildBatchedPadMask, which builds the
+// PREFILL mask ([B,1,L,L], offset 0, right-padded — the training case, already
+// parity-proven by packages/training/tests/native/batch-e2e.test.ts). This module covers the
+// DECODE case for batched serving: left-padded rows sharing one growing
+// [B,H,S,D] KV buffer, where a step forwards N new query tokens at a nonzero
+// offset. (The two builders should be consolidated when batched prefill is
+// wired into the serving path — phase S1a, `02d723a:docs/design/batching.md`.)
+//
+// Why left-padding for decode: with rows right-aligned in the KV buffer, every
+// row's next write lands in the SAME column, so one advancing offset serves
+// the whole batch. Each row's RoPE position is still per-row (handled via the
+// array-offset path, ropeOffsetArr / ops.ropeDynamic — not here); the only
+// thing the mask must encode is (a) causality, (b) each row's left padding,
+// and (c) the sliding window, if any.
+
+import { cacheSignature, isPlainKvCache } from "./capabilities";
+import { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import { createCausalMask } from "../kernels/attention/masks";
+import { runtimeValue } from "../runtime/config";
+import { KVCache } from "./kv";
+import { type AttentionCache, type AttentionRead, type BatchableCache, type Cache, type Mask } from "../contracts/mlx/cache";
+import { AttentionMasks, causalLease, sdpaRead, unmaskedLease, withLease, type MaskLease } from "./attention-read";
+import type { KvStorage } from "./kv-storage";
+
+/** Padding-aware key-validity mask for batched DECODE with left-padded rows.
+ *
+ *  The shared KV buffer has width S (== cache offset once this step's keys are
+ *  written). Row b has `leftPad[b]` padding columns on the left, then its real
+ *  keys. A step forwards N new query tokens whose absolute positions are
+ *  S-N .. S-1 (N == 1 in the steady decode state; N > 1 only for a batched
+ *  prefill chunk).
+ *
+ *  Returns [B, 1, N, S] bool, true where query i of row b may attend to key j:
+ *      allow[b,0,i,j] = (j <= S-N+i)            causal (query i at abs pos S-N+i)
+ *                       AND (j >= leftPad[b])    key past row b's left padding
+ *                       AND (S-N+i < j+window)   within the sliding window*
+ *      (*only when windowSize !== null)
+ *
+ *  Built from createCausalMask(N, S-N, windowSize) (the model's own [N, S] bool
+ *  matrix) broadcast over B, AND-ed with a per-row [B,1,1,S] key-validity mask
+ *  — exactly mirroring buildBatchedPadMask, but at a decode offset with
+ *  left-padding instead of offset-0 right-padding. Caller owns the result. */
+export function buildBatchedDecodeMask(
+  B: number,
+  N: number,
+  S: number,
+  leftPad: number[],
+  windowSize: number | null,
+): MlxArray {
+  // [N, S] causal(+window): queries at abs positions S-N..S-1, keys 0..S-1.
+  const causal2d = createCausalMask(N, S - N, windowSize); // [N, S] bool
+  const causal4d = ops.reshape(causal2d, [1, 1, N, S]);
+  causal2d.dispose();
+
+  // Per-row key validity: keyValid[b,0,0,j] = j >= leftPad[b]. Build as int32
+  // {0,1} then compare to bool, broadcasting over the query axis (i).
+  const keyValidData = new Int32Array(B * S);
+  for (let b = 0; b < B; b++) {
+    const pad = leftPad[b]!;
+    for (let j = 0; j < S; j++) keyValidData[b * S + j] = j >= pad ? 1 : 0;
+  }
+  const keyValidI = MlxArray.fromInt32(keyValidData, [B, 1, 1, S]);
+  const zero = MlxArray.fromInt32(new Int32Array([0]), []);
+  const keyValid = ops.less(zero, keyValidI); // 0 < x → bool, true where x==1
+  zero.dispose();
+  keyValidI.dispose();
+
+  const allow = ops.logicalAnd(causal4d, keyValid); // [1,1,N,S] & [B,1,1,S]
+  causal4d.dispose();
+  keyValid.dispose();
+  return allow; // [B, 1, N, S] bool
+}
+
+/** Cache wrapper for batched DECODE with left-padded rows. Delegates all KV
+ *  storage to a real, shape-[B,H,S,D] cache (KVCache et al. are already
+ *  B-generic), but overrides two things the stock single-sequence path gets
+ *  wrong for a padded batch:
+ *
+ *   - `makeMask` returns the per-row left-padding decode mask
+ *     (buildBatchedDecodeMask) instead of the stock mask, which is EMPTY at
+ *     N==1 and would let every row attend to its own (garbage) left padding.
+ *   - `ropeOffsetArr` exposes each row's REAL position (= pre-write offset −
+ *     leftPad[b]) as an int32 [B] array, so RoPE is applied per-row via the
+ *     array-offset path (mlx fast_rope_dynamic — verified to rotate each row
+ *     by its own offset). The scalar offset would mis-position every padded
+ *     row by its pad amount.
+ *
+ *  `leftPad` is fixed for the batch's lifetime (set when the left-padded decode
+ *  buffer is assembled from per-row prefills). The position array is recomputed
+ *  only when the inner offset advances (once per decode step), so the two reads
+ *  within a step (K before updateAndFetch, Q after — see Attention.forward)
+ *  return the SAME array; the attention graph holds it until eval, so it must
+ *  not be freed mid-step. */
+export class BatchedDecodeMaskCache implements Cache {
+  #ropeArr: MlxArray | null = null;
+  #ropeForOffset = -1;
+
+  /** A decorator, not a storage kind: never claims the inner cache's
+   *  identity (isPlainKvCache(wrapper) must stay false — it is not a
+   *  KVCache), but records what it wraps for diagnostics/persistence. */
+  signature(): string { return `kv:batched-mask:${cacheSignature(this.inner)}`; }
+
+  constructor(
+    private readonly inner: KvStorage,
+    private readonly B: number,
+    private readonly leftPad: number[],
+    private readonly window: number | null,
+  ) {}
+
+  get offset(): number {
+    return this.inner.offset;
+  }
+
+  get ropeOffsetArr(): MlxArray {
+    const off = this.inner.offset;
+    if (this.#ropeArr && this.#ropeForOffset === off) return this.#ropeArr;
+    this.#ropeArr?.dispose(); // prior step's array; its graph is already evaluated
+    const data = new Int32Array(this.B);
+    for (let b = 0; b < this.B; b++) data[b] = off - this.leftPad[b]!;
+    this.#ropeArr = MlxArray.fromInt32(data, [this.B]);
+    this.#ropeForOffset = off;
+    return this.#ropeArr;
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    return this.inner.updateAndFetch(k, v);
+  }
+
+  makeMask(N: number, _windowSize: number | null): Mask {
+    // forwardLayers calls makeMask BEFORE this step's updateAndFetch, so the
+    // KV the attention then fetches spans inner.offset + N keys. Build a fresh
+    // mask each step and hand the model ownership — forwardLayers disposes
+    // mask.arr after building the layer graph (mlx refcounts it into the sdpa
+    // nodes, so post-build dispose is safe), so no view bookkeeping is needed.
+    const S = this.inner.offset + N;
+    return { mode: "array", arr: buildBatchedDecodeMask(this.B, N, S, this.leftPad, this.window) };
+  }
+
+  /** Free the per-step RoPE position array WITHOUT disposing the delegated KV
+   *  (inner). For callers that rebuild a fresh wrapper around a persistent inner
+   *  cache each decode step (the dynamic-B scheduler / parity harness) — the
+   *  inner outlives the wrapper, so wrapper.dispose() would be wrong. */
+  releaseRopeArr(): void {
+    this.#ropeArr?.dispose();
+    this.#ropeArr = null;
+    this.#ropeForOffset = -1;
+  }
+
+  state(): MlxArray[] {
+    return this.inner.state();
+  }
+  isTrimmable(): boolean {
+    return this.inner.isTrimmable();
+  }
+  trim(n: number): void {
+    this.inner.trim(n);
+  }
+  dispose(): void {
+    this.inner.dispose();
+    this.#ropeArr?.dispose();
+    this.#ropeArr = null;
+  }
+}
+
+// --- Dynamic-B cache ops (continuous batching) ---------------------------
+// Ports of mlx-lm's per-cache batch methods (merge / filter), the operations
+// the --batch N scheduler needs that our fixed-B verified forward doesn't have.
+// They run per layer on the raw [.,H,.,D] K/V; the scheduler applies them
+// across layers and wraps the result in a BatchedDecodeMaskCache. Pure array
+// surgery; numerical correctness is gated separately vs an mlx-lm dynamic golden.
+
+/** mlx-lm `merge`: stack N single-row KV slices into one left-padded
+ *  [B,H,Smax,D] batch. Each row is its sequence's KV sliced to its true length
+ *  L_i; we left-pad each by Smax−L_i so all rows align on the right edge (the
+ *  layout decode + BatchedDecodeMaskCache expect). Generalizes
+ *  realBatchedGreedy's assemble to N rows. Caller owns the result; the input
+ *  row arrays are not disposed. */
+export function mergeKVRows(
+  rows: { keys: MlxArray; values: MlxArray }[],
+): { keys: MlxArray; values: MlxArray; leftPad: number[]; width: number } {
+  const lens = rows.map((r) => r.keys.shape[2]!);
+  const width = Math.max(...lens);
+  const leftPad = lens.map((l) => width - l);
+  const leftPadTo = (a: MlxArray, pad: number): MlxArray => {
+    if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]); // fresh full view
+    const [B, H, , D] = a.shape as [number, number, number, number];
+    const z = ops.zeros([B, H, pad, D], a.dtype);
+    const out = ops.concatAxis([z, a], 2);
+    z.dispose();
+    return out;
+  };
+  const ks = rows.map((r, i) => leftPadTo(r.keys, leftPad[i]!));
+  const vs = rows.map((r, i) => leftPadTo(r.values, leftPad[i]!));
+  const keys = ops.concatAxis(ks, 0);
+  const values = ops.concatAxis(vs, 0);
+  for (const a of [...ks, ...vs]) a.dispose();
+  return { keys, values, leftPad, width };
+}
+
+/** mlx-lm `BatchKVCache.extend`: append ONE new right-justified row to an
+ *  existing batched [B,H,S,D] buffer WITHOUT re-extracting the running rows
+ *  (the O(B·S) slice storm mergeKVRows pays on a join). Pad whichever side
+ *  is narrower on the LEFT to the common width, then a single B-axis concat.
+ *  Existing rows' left-pads grow by the batch-side pad (they never shrink —
+ *  mlx-lm extend semantics; the decode mask covers pads either way).
+ *  Caller owns the result; inputs are not disposed. */
+export function extendKVRows(
+  keys: MlxArray, values: MlxArray, leftPad: number[],
+  row: { keys: MlxArray; values: MlxArray },
+): { keys: MlxArray; values: MlxArray; leftPad: number[]; width: number } {
+  const S = keys.shape[2]!;
+  const Lr = row.keys.shape[2]!;
+  const width = Math.max(S, Lr);
+  const padTo = (a: MlxArray, pad: number): MlxArray => {
+    if (pad === 0) return a.slice([0, 0, 0, 0], a.shape as number[]); // fresh full view
+    const [B, H, , D] = a.shape as [number, number, number, number];
+    const z = ops.zeros([B, H, pad, D], a.dtype);
+    const out = ops.concatAxis([z, a], 2);
+    z.dispose();
+    return out;
+  };
+  const bk = padTo(keys, width - S);
+  const bv = padTo(values, width - S);
+  const rk = padTo(row.keys, width - Lr);
+  const rv = padTo(row.values, width - Lr);
+  const outK = ops.concatAxis([bk, rk], 0);
+  const outV = ops.concatAxis([bv, rv], 0);
+  for (const a of [bk, bv, rk, rv]) a.dispose();
+  return {
+    keys: outK,
+    values: outV,
+    leftPad: [...leftPad.map((p) => p + (width - S)), width - Lr],
+    width,
+  };
+}
+
+/** mlx-lm `BatchKVCache.extract` (models/cache.py:1080): pull row `i` out of
+ *  a batched full-attention cache into a fresh SERIAL KVCache — the row slice
+ *  past its left padding, materialized as an OWNED contiguous copy (extraction
+ *  outlives the batch; the source buffers are about to be filter()ed or
+ *  disposed). `leftPad` is the scheduler's per-row padding for this row (full
+ *  layers share one offset; the scheduler owns the pad list). Bit-exact vs a
+ *  solo run by construction: merge/extend/filter are byte-preserving per row
+ *  (`02d723a:tests/parity/batched-decode-parity.test.ts`,
+ *  `02d723a:tests/unit/batched-extract.test.ts`) and this is a pure
+ *  slice+copy of those bytes. Off the hot decode path. */
+export function extractKVRow(cache: KVCache, leftPad: number, i: number): KVCache {
+  if (!cache.keys || !cache.values) throw new Error("extractKVRow: empty cache");
+  const S = cache.offset; // valid columns (buffer may be STEP-padded past it)
+  const cut = (a: MlxArray): MlxArray => {
+    const [, H, , D] = a.shape as [number, number, number, number];
+    const view = a.slice([i, 0, leftPad, 0], [i + 1, H, S, D]);
+    const own = ops.copyOf(view); // TRUE copy: contiguous(view) is a no-op VIEW when already contiguous — pins the source buffer (2026-08-20 DeltaNet conv leak class)
+    view.dispose();
+    return own;
+  };
+  const out = new KVCache();
+  out.restoreState(cut(cache.keys), cut(cache.values), S - leftPad);
+  return out;
+}
+
+/** mlx-lm `filter`: keep only `keep` (sorted row indices) along the batch
+ *  axis — eviction of finished sequences from a batched [B,H,S,D] buffer.
+ *  `trimLeft` removes padding shared by all survivors (the caller updates
+ *  offsets and per-row padding). Caller owns the result; inputs stay owned. */
+export function filterKVRows(
+  keys: MlxArray, values: MlxArray, keep: number[], trimLeft = 0,
+): { keys: MlxArray; values: MlxArray } {
+  const idx = MlxArray.fromInt32(Int32Array.from(keep), [keep.length]);
+  const take = (a: MlxArray) => {
+    const rows = ops.takeAxis(a, idx, 0);
+    if (trimLeft === 0) return rows;
+    const trimmed = rows.slice([0, 0, trimLeft, 0], rows.shape);
+    rows.dispose();
+    return trimmed;
+  };
+  const k = take(keys);
+  const v = take(values);
+  idx.dispose();
+  return { keys: k, values: v };
+}
+
+/** The layout of a running batch's plain full-attention KV: one [B,H,S,D] buffer
+ *  whose rows are right-aligned (left-padded), so every row appends into the same
+ *  column (mlx-lm BatchKVCache). It composes the primitives above: `mergeRows`
+ *  joins one solo row to a running batch (`extendKVRows`, or the O(B·S) re-merge
+ *  with MLX_BUN_BATCH_EXTEND=0), `filterRows` evicts rows and their shared
+ *  padding, `extractRow` pulls a row back out as an owned serial cache. While any
+ *  row is padded the layout gives attention the padding-aware decode mask and
+ *  per-row RoPE positions (BatchedDecodeMaskCache); while none is, it is exactly
+ *  the serial cache, so a batch of equal-length rows runs the serial graph.
+ *  Sources of `mergeRows` are borrowed: a running layout or one serial KVCache,
+ *  then one serial KVCache per joining row.
+ *
+ *  Reads: append to the shared buffer, then the stock fused SDPA. While no row
+ *  is padded, decode passes no mask and a window the fused causal mask, as the
+ *  serial cache does; while one is, both pass the padding-aware mask
+ *  (`buildBatchedDecodeMask`) for their query count, built once per forward
+ *  through `masks`. */
+export class PaddedKVRows implements AttentionCache, BatchableCache {
+  #inner = new KVCache();
+  #leftPad: number[] = [];
+  #padded = false;
+  #view: BatchedDecodeMaskCache | null = null;
+
+  constructor(
+    /** The per-forward masks of the model this layout serves. */
+    readonly masks = new AttentionMasks(),
+  ) {}
+
+  signature(): string { return "kv:padded-rows"; }
+  get batchSize(): number | null { return this.#leftPad.length || null; }
+  get leftPad(): readonly number[] { return this.#leftPad; }
+  get rowOffsets(): readonly number[] { return this.#leftPad.map((pad) => this.#inner.offset - pad); }
+  get offset(): number { return this.#inner.offset; }
+  /** Present only while a row is padded: each row's own position. */
+  get ropeOffsetArr(): MlxArray | undefined { return this.#padded ? this.#step().ropeOffsetArr : undefined; }
+  makeEmptyBatch(): PaddedKVRows { return new PaddedKVRows(this.masks); }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.#read(k, v, this.#mask(1, unmaskedLease));
+  }
+
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.#read(k, v, this.#mask(k.shape[2]!, causalLease));
+  }
+
+  #read(k: MlxArray, v: MlxArray, mask: MaskLease): AttentionRead {
+    return withLease(mask, held => {
+      const [keys, values] = this.#inner.updateAndFetch(k, v);
+      return sdpaRead(keys, values, held);
+    });
+  }
+
+  /** `unpadded` while no row is padded; otherwise the padding-aware mask of an
+   * `N`-position append, before it. */
+  #mask(N: number, unpadded: MaskLease): MaskLease {
+    if (!this.#padded) return unpadded;
+    const S = this.#inner.offset + N;
+    return this.masks.lease("padded-kv", `${N}|${S}|${this.#leftPad.join(",")}`,
+      () => ({ mode: "array", arr: buildBatchedDecodeMask(this.#leftPad.length, N, S, this.#leftPad, null) }));
+  }
+
+  bytesPerToken(): number {
+    return this.#leftPad.length ? this.#inner.bytesPerToken() / this.#leftPad.length : 0;
+  }
+  projectedBytes(tokens: number): number { return this.bytesPerToken() * tokens; }
+
+  #step(): BatchedDecodeMaskCache {
+    return this.#view ??= new BatchedDecodeMaskCache(this.#inner, this.#leftPad.length, [...this.#leftPad], null);
+  }
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] { return this.#inner.updateAndFetch(k, v); }
+  makeMask(N: number, windowSize: number | null): Mask {
+    return this.#padded ? this.#step().makeMask(N, windowSize) : this.#inner.makeMask(N, windowSize);
+  }
+  /** Free this step's RoPE positions after the step's graph is built; the KV stays. */
+  releaseRopeArr(): void { this.#view?.releaseRopeArr(); }
+  state(): MlxArray[] { return this.#inner.state(); }
+  isTrimmable(): boolean { return this.#inner.isTrimmable(); }
+  trim(n: number): void { this.#inner.trim(n); }
+
+  #set(inner: KVCache, leftPad: number[]): void {
+    this.#view?.releaseRopeArr(); this.#view = null;
+    this.#inner = inner; this.#leftPad = leftPad; this.#padded = leftPad.some((pad) => pad !== 0);
+  }
+
+  mergeRows(rows: readonly Cache[]): void {
+    if (this.#leftPad.length) throw new Error("PaddedKVRows.mergeRows needs an empty batch");
+    if (!rows.length) return;
+    const held = new Set<MlxArray>();
+    const own = (array: MlxArray): MlxArray => { held.add(array); return array; };
+    type Part = { keys: MlxArray; values: MlxArray; pads: number[] };
+    const partOf = (cache: Cache): Part => {
+      if (cache instanceof PaddedKVRows) {
+        const [keys, values] = cache.#inner.temporalView();
+        return { keys: own(keys), values: own(values), pads: [...cache.#leftPad] };
+      }
+      if (!isPlainKvCache(cache)) throw new Error(`full-attention batch cannot merge ${cache.signature()}`);
+      const [keys, values] = cache.temporalView();
+      return { keys: own(keys), values: own(values), pads: [0] };
+    };
+    const extend = runtimeValue("MLX_BUN_BATCH_EXTEND") !== "0";
+    try {
+      let acc = partOf(rows[0]!);
+      if (rows.length === 1) {
+        const single = mergeKVRows(acc.pads.map((pad, b) => ({
+          keys: own(acc.keys.slice([b, 0, pad, 0], [b + 1, acc.keys.shape[1]!, acc.keys.shape[2]!, acc.keys.shape[3]!])),
+          values: own(acc.values.slice([b, 0, pad, 0], [b + 1, acc.values.shape[1]!, acc.values.shape[2]!, acc.values.shape[3]!])),
+        })));
+        acc = { keys: own(single.keys), values: own(single.values), pads: single.leftPad };
+      }
+      for (const source of rows.slice(1)) {
+        const row = partOf(source);
+        if (row.pads.length !== 1) throw new Error("PaddedKVRows joins one solo row at a time");
+        let joined: { keys: MlxArray; values: MlxArray; leftPad: number[] };
+        if (extend) {
+          // mlx-lm BatchKVCache.extend: one pad and one concat; existing pads grow, never shrink.
+          joined = extendKVRows(acc.keys, acc.values, acc.pads, row);
+        } else {
+          // The O(B·S) re-merge re-normalizes every row's padding.
+          const [, H, off, D] = acc.keys.shape as [number, number, number, number];
+          const vD = acc.values.shape[3]!;
+          const rowsOut = acc.pads.map((pad, b) => ({
+            keys: own(acc.keys.slice([b, 0, pad, 0], [b + 1, H, off, D])),
+            values: own(acc.values.slice([b, 0, pad, 0], [b + 1, H, off, vD])),
+          }));
+          joined = mergeKVRows([...rowsOut, row]);
+        }
+        acc = { keys: own(joined.keys), values: own(joined.values), pads: joined.leftPad };
+      }
+      held.delete(acc.keys); held.delete(acc.values);
+      const cache = new KVCache();
+      cache.restoreState(acc.keys, acc.values, acc.keys.shape[2]!);
+      this.#set(cache, acc.pads);
+    } finally { for (const array of held) array.dispose(); }
+  }
+
+  extractRow(row: number): Cache {
+    return this.#inner.keys ? extractKVRow(this.#inner, this.#leftPad[row]!, row) : new KVCache();
+  }
+
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    const kept = keep.map((row) => this.#leftPad[row]!);
+    // mlx-lm BatchKVCache.filter removes the padding every survivor shares.
+    const shared = Math.min(...kept);
+    const [keys, values] = this.#inner.temporalView();
+    let filtered: { keys: MlxArray; values: MlxArray };
+    try { filtered = filterKVRows(keys, values, [...keep], shared); }
+    finally { keys.dispose(); values.dispose(); }
+    const cache = new KVCache();
+    cache.restoreState(filtered.keys, filtered.values, this.#inner.offset - shared);
+    this.#inner.dispose();
+    this.#set(cache, kept.map((pad) => pad - shared));
+  }
+
+  dispose(): void {
+    this.#inner.dispose();
+    this.#set(new KVCache(), []);
+    this.masks.clear();
+  }
+}