@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/kv-tensor-rows.ts

@@ -0,0 +1,212 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import { materializeCopy } from "@mlx-bun/mlx/materialize";
+import { FullPrefillPadding } from "./full-prefill-padding";
+import type { Mask, PrefillPadding } from "../contracts/mlx/cache";
+import { cleanupFailure, disposeResources } from "../runtime/resources";
+
+/** Borrowed tensor planes with shared token positions. Plane shapes may differ
+ * in head width and dtype. Every plane uses [row, head, token, field]. */
+export interface KvTensorRowView {
+  readonly planes: readonly MlxArray[];
+  readonly rowOffsets: readonly number[];
+  readonly leftPad: readonly number[];
+}
+
+/** Position and membership operations shared by packed KV representations.
+ * A codec supplies tensor planes; this owner never encodes or decodes them. */
+export class KvTensorRows implements KvTensorRowView {
+  planes: MlxArray[] = [];
+  rowOffsets: number[] = [];
+  leftPad: number[] = [];
+  #rope: MlxArray | undefined;
+  #beforeRound: number[] | undefined;
+  readonly #padding = new FullPrefillPadding();
+
+  restorePrefillEnds(ends: readonly number[] | undefined): void { this.#padding.restoreEnds(ends); }
+  preparePrefill(padding: PrefillPadding): void {
+    this.#padding.prepare(this, padding); this.#positionsChanged();
+  }
+  finalizePrefill(): void {
+    const next = this.#padding.finalize(this.planes, this);
+    if (next) this.#replace(next, this.rowOffsets, this.leftPad);
+  }
+
+  /** Snapshot logical validity without changing physical row alignment. */
+  captureDonorValidity(): Pick<import("../contracts/mlx/cache").KvDonorRows, "offsets" | "starts" | "ends"> {
+    const offsets = this.rowOffsets.map((offset, row) => Math.max(0, this.#padding.validOffset(offset, row)));
+    return { offsets, starts: [...this.leftPad], ends: offsets.map((offset, row) => offset + this.leftPad[row]!) };
+  }
+  get offset(): number { return Math.max(0, ...this.#ends()); }
+  get batchSize(): number | null { return this.rowOffsets.length || null; }
+  #ends(): number[] { return this.rowOffsets.map((offset, row) => offset + this.leftPad[row]!); }
+  get ropeOffsetArr(): MlxArray | undefined {
+    if (this.leftPad.every(pad => pad === 0) && this.rowOffsets.every(offset => offset === this.rowOffsets[0])) return undefined;
+    return this.#rope ??= ops.fromInt32(this.rowOffsets, [this.rowOffsets.length]);
+  }
+  #positionsChanged(): void { this.#rope?.dispose(); this.#rope = undefined; }
+  bytesPerToken(): number {
+    return this.planes.length && this.rowOffsets.length ? this.planes.reduce((bytes, plane) => bytes + plane.nbytes, 0) /
+      (this.rowOffsets.length * this.planes[0]!.shape[2]!) : 0;
+  }
+  trim(count: number): void {
+    this.rowOffsets = this.rowOffsets.map(offset => offset - count); this.#positionsChanged();
+  }
+  specRoundBegin(): void { this.#beforeRound = [...this.rowOffsets]; }
+  specRoundCommit(): void { this.#beforeRound = undefined; }
+  specRoundRollback(keep: number | readonly number[]): void {
+    this.rowOffsets = this.#beforeRound!.map((offset, row) => offset + (typeof keep === "number" ? keep : keep[row]!));
+    this.#beforeRound = undefined; this.#positionsChanged();
+  }
+  makeMask(N: number, window: number | null): Mask {
+    const B = this.rowOffsets.length, ends = this.#ends(), S = this.offset + N;
+    if (window === null && this.leftPad.every(pad => pad === 0) && ends.every(end => end === ends[0]))
+      return { mode: N === 1 ? "" : "causal", arr: null };
+    using starts = ops.fromInt32(ends, [B, 1, 1, 1]);
+    using relative = ops.arange(0, N, 1, Dtype.int32);
+    using queries = ops.reshape(relative, [1, 1, N, 1]);
+    using position = ops.add(starts, queries);
+    using indices = ops.arange(0, S, 1, Dtype.int32);
+    using keys = ops.reshape(indices, [1, 1, 1, S]);
+    using pads = ops.fromInt32(this.leftPad, [B, 1, 1, 1]);
+    using causal = ops.lessEqual(keys, position);
+    using valid = ops.greaterEqual(keys, pads);
+    const mask = ops.logicalAnd(causal, valid);
+    if (window === null) return { mode: "array", arr: mask };
+    using windowSize = ops.fromInt32([window], []);
+    using limit = ops.add(keys, windowSize);
+    using within = ops.less(position, limit);
+    try { return { mode: "array", arr: ops.logicalAnd(mask, within) }; }
+    finally { mask.dispose(); }
+  }
+  #replace(planes: MlxArray[], offsets: number[], pads: number[]): void {
+    const previous = this.planes;
+    this.planes = planes; this.rowOffsets = offsets; this.leftPad = pads;
+    this.#positionsChanged(); disposeResources(previous);
+  }
+  append(inputs: readonly MlxArray[]): void {
+    const B = this.rowOffsets.length, N = inputs[0]!.shape[2]!, ends = this.#ends();
+    const needed = Math.max(...ends) + N, capacity = Math.ceil(needed / 256) * 256;
+    const common = ends.every(end => end === ends[0]);
+    using positions = common ? null : ops.fromInt32(ends.flatMap(end =>
+      Array.from({ length: N }, (_, position) => end + position)), [B, 1, N, 1]);
+    const next: MlxArray[] = [];
+    try {
+      for (const [index, input] of inputs.entries()) {
+        const old = this.planes[index];
+        let base: MlxArray;
+        if (!old) base = ops.zeros([B, input.shape[1]!, capacity, input.shape[3]!], input.dtype);
+        else if (old.shape[2]! >= needed) base = old;
+        else {
+          using padding = ops.zeros([B, old.shape[1]!, capacity - old.shape[2]!, old.shape[3]!], old.dtype);
+          base = ops.concatAxis([old, padding], 2);
+        }
+        try {
+          const updated = common
+            ? ops.sliceUpdate(base, input, [0, 0, ends[0]!, 0], [B, base.shape[1]!, ends[0]! + N, base.shape[3]!])
+            : ops.putAlongAxis(base, positions!, input, 2);
+          next.push(updated);
+        } finally { if (base !== old) base.dispose(); }
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    this.#replace(next, this.rowOffsets.map(offset => offset + N), this.leftPad);
+  }
+  mergeRows(rows: readonly KvTensorRowView[]): void {
+    const next: MlxArray[] = [], held: MlxArray[] = [];
+    let offsets: number[], pads: number[];
+    try {
+      if (rows.length === 1) {
+        const row = rows[0]!;
+        for (const plane of row.planes) next.push(ops.copyOf(plane));
+        offsets = [...row.rowOffsets]; pads = [...row.leftPad];
+      } else {
+        const widths = rows.map(row => Math.max(0, ...row.rowOffsets.map((offset, i) => offset + row.leftPad[i]!)));
+        const width = Math.max(0, ...widths);
+        if (width) {
+          const prototype = rows.find((_, index) => widths[index]! > 0)!;
+          for (const [field, geometry] of prototype.planes.entries()) {
+            const pieces = rows.map((row, index) => {
+              const active = widths[index]!;
+              if (!active) {
+                const empty = ops.zeros([row.rowOffsets.length, geometry.shape[1]!, width, geometry.shape[3]!], geometry.dtype);
+                held.push(empty); return empty;
+              }
+              const plane = row.planes[field]!;
+              const view = plane.slice([0, 0, 0, 0], [row.rowOffsets.length, plane.shape[1]!, active, plane.shape[3]!]);
+              held.push(view);
+              if (active === width) return view;
+              using padding = ops.zeros([row.rowOffsets.length, plane.shape[1]!, width - active, plane.shape[3]!], plane.dtype);
+              const padded = ops.concatAxis([padding, view], 2); held.push(padded); return padded;
+            });
+            next.push(ops.concatAxis(pieces, 0));
+          }
+        }
+        offsets = rows.flatMap(row => [...row.rowOffsets]);
+        pads = rows.flatMap((row, index) => row.leftPad.map(pad => pad + width - widths[index]!));
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    finally { disposeResources(held); }
+    this.#replace(next, offsets, pads);
+  }
+  /** Preserve caller-owned physical positions when changing representation.
+   * Live tensor bytes are copied; padding never participates in a codec. */
+  alignRows(leftPad: readonly number[]): void {
+    if (leftPad.every((pad, row) => pad === this.leftPad[row])) return;
+    const width = Math.max(0, ...this.rowOffsets.map((offset, row) => offset + leftPad[row]!));
+    const rows: MlxArray[][] = [], planes: MlxArray[] = [];
+    try {
+      for (let row = 0; row < this.rowOffsets.length; row++) rows.push(this.extractRow(row));
+      for (let field = 0; field < this.planes.length; field++) {
+        const parts: MlxArray[] = [];
+        try {
+          for (let row = 0; row < rows.length; row++) {
+            const geometry = this.planes[field]!, live = rows[row]![field];
+            using before = ops.zeros([1, geometry.shape[1]!, leftPad[row]!, geometry.shape[3]!], geometry.dtype);
+            using after = ops.zeros([1, geometry.shape[1]!, width - leftPad[row]! - this.rowOffsets[row]!, geometry.shape[3]!], geometry.dtype);
+            parts.push(ops.concatAxis(live ? [before, live, after] : [before, after], 2));
+          }
+          planes.push(ops.concatAxis(parts, 0));
+        } finally { disposeResources(parts); }
+      }
+    } catch (error) { return cleanupFailure(error, () => disposeResources(planes)); }
+    finally { disposeResources(rows.flat()); }
+    this.#replace(planes, [...this.rowOffsets], [...leftPad]);
+  }
+  filterRows(keep: readonly number[]): void {
+    if (!keep.length) { this.dispose(); return; }
+    const next: MlxArray[] = [];
+    using indices = this.planes.length ? ops.fromInt32([...keep], [keep.length]) : null;
+    const pads = keep.map(row => this.leftPad[row]!);
+    const sharedPadding = Math.min(...pads);
+    try {
+      for (const plane of this.planes) {
+        const selected = ops.takeAxis(plane, indices!, 0);
+        if (!sharedPadding) next.push(selected);
+        else {
+          try { next.push(selected.slice([0, 0, sharedPadding, 0], [...selected.shape])); }
+          finally { selected.dispose(); }
+        }
+      }
+    }
+    catch (error) { return cleanupFailure(error, () => disposeResources(next)); }
+    this.#replace(next, keep.map(row => this.rowOffsets[row]!), pads.map(pad => pad - sharedPadding));
+    this.#padding.filter(keep);
+  }
+  extractRow(row: number): MlxArray[] {
+    const start = this.leftPad[row]!, count = this.rowOffsets[row]!, result: MlxArray[] = [];
+    if (!count) return result;
+    try {
+      for (const plane of this.planes) {
+        using view = plane.slice([row, 0, start, 0], [row + 1, plane.shape[1]!, start + count, plane.shape[3]!]);
+        result.push(materializeCopy(view));
+      }
+      return result;
+    } catch (error) { return cleanupFailure(error, () => disposeResources(result)); }
+  }
+  dispose(): void {
+    const previous = this.planes;
+    this.planes = []; this.rowOffsets = []; this.leftPad = []; this.#beforeRound = undefined; this.#padding.clear();
+    this.#positionsChanged(); disposeResources(previous);
+  }
+}

package/src/state/kv.ts

@@ -0,0 +1,324 @@
+import { MlxArray } from "@mlx-bun/mlx/array";
+import { Dtype } from "@mlx-bun/mlx/ffi";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { AttentionRead,BidirectionalAttentionCache,BlockAttentionCache,Cache,CommittedAttentionCache,CompiledDecodeCache,DecodeSlot,DecodeStepInputs,DecodeStepPlan,DecodeTrace,EncoderAttentionCache,KvDonorRows,Mask } from "../contracts/mlx/cache";
+import { bidirMask, createCausalMask } from "../kernels/attention/masks";
+import { QuantizedKVCache } from "./quantized-kv";
+import { AttentionMasks, causalLease, encoderWindowMask, sdpaRead, unmaskedLease, withLease, type DenseRead } from "./attention-read";
+import type { AffineKernelSet } from "./affine-attention";
+import { uncomposedTurboQuantKernels, type TurboQuantKernels } from "./turboquant-attention";
+
+export type { AffineKernelSet };
+
+
+/** KV cache — port of mlx-lm cache.py KVCache: preallocated in steps of
+ *  256 along the sequence axis, updated in place via slice_update.
+ *
+ *  Reads: each named read appends through `updateAndFetch` and attends the
+ *  fetched prefix with the stock fused SDPA, the call every graph makes after
+ *  `updateAndFetch` today, a window through `windowRead` (that SDPA unless the
+ *  composition names a read computing its bits otherwise). Decode passes no
+ *  mask and a window the fused causal mask; at one query the two compute the
+ *  same bits, so a one-row window needs no case of its own. This cache has no
+ *  sliding window (graphs compose `RotatingKVCache` for sliding layers), so
+ *  only the bidirectional and encoder reads build a mask array. */
+export class KVCache implements CompiledDecodeCache, CommittedAttentionCache, BidirectionalAttentionCache, EncoderAttentionCache, BlockAttentionCache {
+  static readonly STEP = 256;
+  /** Set only by compiled-decode trace adapters (see Cache). */
+  readonly ropeOffsetArr?: MlxArray;
+  keys: MlxArray | null = null;
+  values: MlxArray | null = null;
+  offset = 0;
+
+  constructor(
+    /** The per-forward masks of the model this cache belongs to. Row layouts
+     * made from this cache (`state/layout`) share them. */
+    readonly masks = new AttentionMasks(),
+    /** The TurboQuant read of the graph that built this cache: a TurboQuant
+     * conversion of this storage (`TurboQuantKVCache.fromKVCache`, the delayed
+     * TurboQuant rows) attends with it, and copies of this cache carry it.
+     * Qwen 3.5's `makeCache` passes `denseTurboQuantKernels`, Gemma 4's and
+     * MiniCPM5's `deferredTurboQuantKernels`; a graph that reads TurboQuant
+     * storage only through the deprecated fetches passes none. */
+    readonly turboQuantKernels: TurboQuantKernels = uncomposedTurboQuantKernels,
+    /** The read of a causal window (`appendWindow`, `appendCommitted`): the
+     * stock fused SDPA, or the read the composition names
+     * (`foldedWindowRead`, bit for bit the same). Row layouts made from this
+     * storage and its clones read their windows with it. */
+    readonly windowRead: DenseRead = sdpaRead,
+  ) {}
+
+  signature(): string { return "kv:plain"; }
+
+  appendDecode(k: MlxArray, v: MlxArray): AttentionRead {
+    const [keys, values] = this.updateAndFetch(k, v);
+    return sdpaRead(keys, values, unmaskedLease);
+  }
+
+  appendWindow(k: MlxArray, v: MlxArray): AttentionRead {
+    const [keys, values] = this.updateAndFetch(k, v);
+    return this.windowRead(keys, values, causalLease);
+  }
+
+  /** A bf16 committed span reads exactly as a window does. */
+  appendCommitted(k: MlxArray, v: MlxArray): AttentionRead {
+    return this.appendWindow(k, v);
+  }
+
+  /** Gemma 4's vision prefill read: the causal mask OR'd with the image block
+   * overlay (`bidirMask`), then the fused SDPA with that array mask. One mask
+   * per forward: every layer passes the same `bidirectional` tensor. */
+  appendBidirectional(k: MlxArray, v: MlxArray, bidirectional: MlxArray): AttentionRead {
+    const L = k.shape[2]!;
+    const mask = this.masks.lease("bidirectional:none", `${L}`, () => bidirMask(L, null, bidirectional), bidirectional);
+    return withLease(mask, held => {
+      const [keys, values] = this.updateAndFetch(k, v);
+      return sdpaRead(keys, values, held);
+    });
+  }
+
+  /** DiffusionGemma's vision encoder read on a full-attention layer: the
+   * append `updateAndFetch` makes, then the fused SDPA under the encoder
+   * window's mask (`encoderWindowMask`), one per forward for every layer. */
+  appendEncoderWindow(k: MlxArray, v: MlxArray, bidirectional: MlxArray): AttentionRead {
+    return withLease(encoderWindowMask(this.masks, k.shape[2]!, bidirectional), mask => {
+      const [keys, values] = this.updateAndFetch(k, v);
+      return sdpaRead(keys, values, mask);
+    });
+  }
+
+  /** DiffusionGemma's canvas read on a full-attention layer: the stored
+   * positions in order, then the block's keys and values, with no mask. */
+  readBlock(k: MlxArray, v: MlxArray): AttentionRead {
+    const [storedK, storedV] = this.temporalView();
+    let keys: MlxArray, values: MlxArray;
+    try {
+      keys = ops.concatAxis([storedK, k], 2);
+      try { values = ops.concatAxis([storedV, v], 2); }
+      catch (error) { keys.dispose(); throw error; }
+    } finally { storedK.dispose(); storedV.dispose(); }
+    return sdpaRead(keys, values, unmaskedLease);
+  }
+
+  bytesPerToken(): number {
+    const steps = this.keys?.shape[2] ?? 0;
+    return steps > 0 ? (this.keys!.nbytes + this.values!.nbytes) / steps : 0;
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const prev = this.offset;
+    const L = k.shape[2]!;
+    if (!this.keys || prev + L > this.keys.shape[2]!) {
+      const [B, H, , D] = k.shape as [number, number, number, number];
+      const vD = v.shape[3]!;
+      const nSteps = Math.floor((KVCache.STEP + L - 1) / KVCache.STEP);
+      const newK = ops.zeros([B, H, nSteps * KVCache.STEP, D], k.dtype);
+      const newV = ops.zeros([B, H, nSteps * KVCache.STEP, vD], v.dtype);
+      if (this.keys && this.values) {
+        let oldK = this.keys;
+        let oldV = this.values;
+        if (prev % KVCache.STEP !== 0) {
+          const trimK = oldK.slice([0, 0, 0, 0], [B, H, prev, D]);
+          const trimV = oldV.slice([0, 0, 0, 0], [B, H, prev, vD]);
+          oldK.dispose();
+          oldV.dispose();
+          oldK = trimK;
+          oldV = trimV;
+        }
+        this.keys = ops.concatAxis([oldK, newK], 2);
+        this.values = ops.concatAxis([oldV, newV], 2);
+        for (const a of [oldK, oldV, newK, newV]) a.dispose();
+      } else {
+        this.keys = newK;
+        this.values = newV;
+      }
+    }
+
+    this.offset += L;
+    const [B, H, S, D] = this.keys!.shape as [number, number, number, number];
+    const vD = this.values!.shape[3]!;
+    const k2 = ops.sliceUpdate(this.keys!, k, [0, 0, prev, 0], [B, H, this.offset, D]);
+    const v2 = ops.sliceUpdate(this.values!, v, [0, 0, prev, 0], [B, H, this.offset, vD]);
+    this.keys!.dispose();
+    this.values!.dispose();
+    this.keys = k2;
+    this.values = v2;
+    return [
+      this.keys.slice([0, 0, 0, 0], [B, H, this.offset, D]),
+      this.values.slice([0, 0, 0, 0], [B, H, this.offset, vD]),
+    ];
+  }
+
+  makeMask(N: number, windowSize: number | null): Mask {
+    if (N === 1) return { mode: "", arr: null };
+    // Windowless multi-token chunks get the STRING "causal" at ANY offset —
+    // mlx-lm cache.py:114-125 create_attention_mask does exactly this, and
+    // mx.fast SDPA aligns the string bottom-right, i.e. the same values as
+    // our materialized matrix. Shipping the array instead (the old
+    // offset>0 branch) forced a materialize-and-add mask path: a bool
+    // [N, offset+N] built PER CHUNK and read by every full-attention layer
+    // — the e4b 16k prefill gap vs mlx-lm (2026-07-06b, hd=512 layers
+    // where the fallback dispatch can't block-skip an array mask).
+    if (windowSize === null) return { mode: "causal", arr: null };
+    if (this.offset === 0 && N <= windowSize) return { mode: "causal", arr: null };
+    return { mode: "array", arr: createCausalMask(N, this.offset, windowSize) };
+  }
+
+  /** Arrays to eval to materialize cache state (prefill chunk boundary). */
+  state(): MlxArray[] {
+    return this.keys && this.values ? [this.keys, this.values] : [];
+  }
+
+  isTrimmable(): boolean {
+    return true;
+  }
+
+  trim(n: number): void {
+    this.offset = Math.max(0, this.offset - n);
+  }
+
+  /** Chronological (K, V) view sliced to offset (drafter donor read). */
+  captureDonorRows(): KvDonorRows {
+    const [keys, values] = this.temporalView();
+    const B = keys.shape[0]!, width = keys.shape[2]!;
+    return { keys, values, offsets: Array(B).fill(this.offset),
+      starts: Array(B).fill(0), ends: Array(B).fill(width) };
+  }
+
+  temporalView(): [MlxArray, MlxArray] {
+    if (!this.keys || !this.values) throw new Error("cache is empty");
+    const [B, H, , D] = this.keys.shape as [number, number, number, number];
+    const vD = this.values.shape[3]!;
+    return [
+      this.keys.slice([0, 0, 0, 0], [B, H, this.offset, D]),
+      this.values.slice([0, 0, 0, 0], [B, H, this.offset, vD]),
+    ];
+  }
+
+  /** Compiled decode: host-side bookkeeping for a 1-token step — the
+   *  growth/trim half of updateAndFetch (L=1), without the write. */
+  prepareDecodeStep(): DecodeStepPlan {
+    const prev = this.offset;
+    if (!this.keys || !this.values) throw new Error("compiled decode on an empty cache");
+    if (prev + 1 > this.keys.shape[2]!) {
+      const [B, H, , D] = this.keys.shape as [number, number, number, number];
+      const vD = this.values.shape[3]!;
+      let oldK = this.keys;
+      let oldV = this.values;
+      if (prev % KVCache.STEP !== 0) {
+        const trimK = oldK.slice([0, 0, 0, 0], [B, H, prev, D]);
+        const trimV = oldV.slice([0, 0, 0, 0], [B, H, prev, vD]);
+        oldK.dispose();
+        oldV.dispose();
+        oldK = trimK;
+        oldV = trimV;
+      }
+      const newK = ops.zeros([B, H, KVCache.STEP, D], oldK.dtype);
+      const newV = ops.zeros([B, H, KVCache.STEP, vD], oldV.dtype);
+      this.keys = ops.concatAxis([oldK, newK], 2);
+      this.values = ops.concatAxis([oldV, newV], 2);
+      for (const a of [oldK, oldV, newK, newV]) a.dispose();
+    }
+    return { fetch: "concat", writePos: prev, activeLen: prev };
+  }
+
+  decodePhase(): DecodeStepPlan["fetch"] { return "concat"; }
+
+  decodeSlot(_plan: DecodeStepPlan): DecodeSlot { return concatDecodeSlot(this.offset); }
+
+  decodeInputs(plan: DecodeStepPlan, step: DecodeStepInputs): MlxArray[] {
+    return [step.activeView(this.keys!, plan.activeLen), step.activeView(this.values!, plan.activeLen)];
+  }
+
+  commitDecodeStep(_slot: DecodeSlot, outputs: MlxArray[]): MlxArray[] {
+    return this.writeDecodeStep(outputs[0]!, outputs[1]!);
+  }
+
+  /** Compiled decode: the write half (same sliceUpdate as updateAndFetch).
+   *  Takes ownership of kNew/vNew; returns the arrays to async-eval with
+   *  the step (the updated buffers). */
+  writeDecodeStep(kNew: MlxArray, vNew: MlxArray): MlxArray[] {
+    const prev = this.offset;
+    const [B, H, , D] = this.keys!.shape as [number, number, number, number];
+    const vD = this.values!.shape[3]!;
+    const k2 = ops.sliceUpdate(this.keys!, kNew, [0, 0, prev, 0], [B, H, prev + 1, D]);
+    const v2 = ops.sliceUpdate(this.values!, vNew, [0, 0, prev, 0], [B, H, prev + 1, vD]);
+    this.keys!.dispose();
+    this.values!.dispose();
+    kNew.dispose();
+    vNew.dispose();
+    this.keys = k2;
+    this.values = v2;
+    this.offset = prev + 1;
+    return [k2, v2];
+  }
+
+  /** Adopt persisted state (takes ownership of the arrays). */
+  restoreState(keys: MlxArray, values: MlxArray, offset: number): void {
+    this.dispose();
+    this.keys = keys;
+    this.values = values;
+    this.offset = offset;
+  }
+
+  /** Port of mlx-lm KVCache.to_quantized: quantize the whole buffer
+   *  (padding included — it's overwritten before being read). The converted
+   *  cache attends with the kernel set `kernels` builds for its bits, group
+   *  size and this cache's activation dtype. */
+  toQuantized(groupSize: number, bits: number, kernels: AffineKernelSet): QuantizedKVCache {
+    const q = new QuantizedKVCache(groupSize, bits, kernels(bits, groupSize, this.keys?.dtype ?? Dtype.bfloat16), this.masks);
+    q.offset = this.offset;
+    if (this.keys && this.values) {
+      q.keys = ops.quantize(this.keys, groupSize, bits);
+      q.values = ops.quantize(this.values, groupSize, bits);
+    }
+    this.dispose();
+    return q;
+  }
+
+  dispose(): void {
+    this.keys?.dispose();
+    this.values?.dispose();
+    this.keys = this.values = null;
+    this.offset = 0;
+    this.masks.clear();
+  }
+}
+
+/** The growing plain cache inside a compiled decode trace (also a rotating
+ *  cache before its window fills): fetch = concat(active prefix, new row); the
+ *  new row is a closure output for the write after the step. `offset` is set for
+ *  completeness only: at L=1 every makeMask returns mode "" and the RoPE offset
+ *  rides `ropeOffsetArr`, so the baked value cannot leak into the graph. */
+class TracedConcatKVCache extends KVCache implements DecodeTrace {
+  override readonly ropeOffsetArr: MlxArray;
+  outs: MlxArray[] = [];
+  constructor(
+    offset: number,
+    readonly activeK: MlxArray,
+    readonly activeV: MlxArray,
+    ropeOffsetArr: MlxArray,
+  ) {
+    super();
+    this.offset = offset;
+    this.ropeOffsetArr = ropeOffsetArr;
+  }
+
+  override updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    // same-shape reshape: an owned alias that survives the caller's
+    // dispose of k/v, so the row can be a closure output
+    this.outs = [ops.reshape(k, k.shape), ops.reshape(v, v.shape)];
+    return [
+      ops.concatAxis([this.activeK, k], 2),
+      ops.concatAxis([this.activeV, v], 2),
+    ];
+  }
+}
+
+/** The concat-fetch slot of a plain cache at `offset` (KVCache, and RotatingKVCache before the window fills). */
+export function concatDecodeSlot(offset: number): DecodeSlot {
+  return {
+    key: "p-cat:0:0", fetch: "concat", inputs: 2, outputs: 2,
+    trace: (inputs, ropeOffset) => new TracedConcatKVCache(offset, inputs[0]!, inputs[1]!, ropeOffset),
+  };
+}

package/src/state/kv4-head256-cache.ts

@@ -0,0 +1,46 @@
+import { KV4_DECODE_BLOCK } from "../kernels/attention/kvq4-decode";
+import { kv4DecodeAttention } from "../layers/kv4-decode-attention";
+import { foldedQuantizedSdpa } from "../layers/quantized-attention";
+import { unfusedKernel, type AffineKernels } from "./affine-attention";
+import { AttentionMasks } from "./attention-read";
+import { QuantizedKVCache } from "./quantized-kv";
+
+/** The kernels of the Qwen3.8-27B Trellis compositions' full-attention layers
+ * for one sequence: the fused one-row KV4 decode kernel (`kv4DecodeAttention`,
+ * each packed key and value row read once for every query head of its KV head,
+ * `decodeBlock` tokens per threadgroup: the M4 Pro's 512 by default, the M1
+ * Max's 128 as its composition names) and the folded-GQA window kernel
+ * (`foldedQuantizedSdpa`, one quantized matmul per KV head for its query rows,
+ * causal over the last L positions). Neither reads the mask, and both take one
+ * row (B = 1), head dim 256 and an even GQA factor. `multiRow` is the
+ * composition's affine set, which a row layout reads with whenever it holds
+ * more than one row or one behind left padding: the old graph served several
+ * sequences on the generic path. A masked window on serial storage takes the
+ * unfused port, which this storage never needs (its mask is plain causal). */
+export function kv4Head256Kernels(multiRow: AffineKernels, decodeBlock = KV4_DECODE_BLOCK): AffineKernels {
+  return {
+    decode: (q, keys, values, scale) => kv4DecodeAttention(q, keys, values, scale, 64, 4, decodeBlock),
+    window: (q, keys, values, scale) => foldedQuantizedSdpa(q, keys, values, scale, 64, 4),
+    maskedWindow: unfusedKernel(64, 4),
+    multiRow,
+  };
+}
+
+/** The 4-bit group-64 head-dim-256 full-attention storage the M4 Pro Qwen3.8
+ * composition reads, quantized from its first append: `QuantizedKVCache`
+ * storage with `kv4Head256Kernels`. The served lego is the bf16-first
+ * `Bf16FirstKv4Head256Cache` (`state/bf16-first-kv.ts`), whose converted
+ * storage reads with these kernels; this cache is the quantized-from-token-zero
+ * composition. The committed read is `QuantizedKVCache`'s.
+ *
+ * Each read is bit-identical to the kernel the old graph injected for one
+ * sequence; neither kernel is bit-identical to the unfused port (f32 online
+ * softmax, regrouped rows: bf16-level). A one-row window takes the folded
+ * kernel, where the old graph's one-row plan took the KV4 decode kernel. Row
+ * layouts made from this cache keep its kernels, `multiRow` included. */
+export class Kv4Head256Cache extends QuantizedKVCache {
+  /** `multiRow`: the composition's affine kernels for 4-bit group 64. */
+  constructor(multiRow: AffineKernels, masks = new AttentionMasks()) {
+    super(64, 4, kv4Head256Kernels(multiRow), masks);
+  }
+}