mlx-bun 0.3.0 → 0.5.0

package/src/model/gemma4-base.ts

@@ -7,15 +7,19 @@
 // DecoderLayer, Gemma4Model) stays in gemma4.ts and becomes the
 // per-model generated code in Phase C.
 
+import { decodedKvDonorAttention } from "./decoded-kv-donor";
 import type { ModelConfig } from "../config";
 import { quantFor } from "../config";
 import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
-import { Dtype } from "../mlx/ffi";
+import { Dtype, deviceArchitecture } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { quantizedMatmulRows } from "../mlx/quantized-rows";
 import { expertOffloadArray } from "../expert-offload";
 import * as tq from "../mlx/turboquant-ops";
+import type { KvCodec } from "../backends/mlx/kv-codec";
+import { TurboQuantCodec, turboQuantFusedDecode, disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
+export { disposeTurboQuant, type TurboQuantTensor } from "./turboquant-codec";
 import { runtimeValue } from "../runtime-config";
 
 export type MaskMode = "" | "causal";
@@ -249,9 +253,9 @@ export class QuantizedEmbedding {
 
 /** Fetched KV handed from a donor layer to its sharers (and to the
  *  speculative drafter). Owned by forwardLayers for the pass.
- *  `offsetArr` is set only under compiled decode (trace adapters): the
- *  RoPE offset as an array VALUE so the compiled graph replays at any
- *  position (see src/model/compiled-decode.ts). */
+ *  `offsetArr` is an owned pre-write position snapshot. Donor and sharing
+ *  attention layers keep it until the fetched KV is released, even if the
+ *  cache advances or replaces its borrowed row-position array. */
 export type SharedKv =
   | { kind: "plain"; keys: MlxArray; values: MlxArray; offset: number;
       offsetArr?: MlxArray;
@@ -260,11 +264,90 @@ export type SharedKv =
        *  (tq.unrotateValues) after sdpa. Attention is linear in V, so the
        *  result is the same decode, just transformed once per query row
        *  instead of once per cached token. */
-      vRotated?: boolean }
+      vRotated?: boolean;
+      restoreValues?: (output: MlxArray) => MlxArray }
   | { kind: "quant"; keys: ops.QuantizedTensor; values: ops.QuantizedTensor;
-      offset: number; groupSize: number; bits: number; offsetArr?: MlxArray };
+      offset: number; groupSize: number; bits: number; offsetArr?: MlxArray }
+  | { kind: "view"; attention: KvAttentionView; offset: number; offsetArr?: MlxArray };
+
+/** Retain an independent array handle for positions used across cache writes.
+ * The contiguous view preserves the integer values and shares their storage;
+ * callers release it with the fetched attention state. */
+export function captureRopeOffsets(cache: Pick<Cache, "ropeOffsetArr">): MlxArray | undefined {
+  const offsets = cache.ropeOffsetArr;
+  return offsets ? ops.contiguous(offsets) : undefined;
+}
+
+/** Quantized attention consumes a numerical storage port, independent of the
+ * concrete layout used for row positions, retention and persistence. */
+export interface QuantizedAttentionState {
+  readonly groupSize: number;
+  readonly bits: number;
+  updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor];
+}
+
+/** Values may stay in the codec's rotated domain until after attention. */
+export interface RotatedValueAttentionState {
+  /** Capture a fetched value domain independently of later row changes. */
+  captureValueTransform?(): (output: MlxArray) => MlxArray;
+  updateAndFetchDeferredV(k: MlxArray, v: MlxArray): [MlxArray, MlxArray];
+}
+
+/** A captured numerical view supports multiple queries without appending KV
+ * again. It owns its tensor handles independently of later cache membership or
+ * precision changes. Query/mask inputs are borrowed; outputs are owned. */
+export interface KvAttentionView {
+  /** Qualified committed spans retain the single-position reduction order. */
+  attend(q: MlxArray, scale: number, mask: Mask, independentPositions?: boolean): MlxArray;
+  dispose(): void;
+}
+
+/** Storage appends once and hands every attention consumer the same view. */
+export interface KvAttentionState {
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView;
+}
+
+/** Owned read-only planes with validity separate from physical padding. */
+export interface KvDonorRows {
+  readonly keys: MlxArray;
+  readonly values: MlxArray;
+  readonly offsets: readonly number[];
+  readonly starts: readonly number[];
+  readonly ends: readonly number[];
+}
+
+/** Captured Q-only attention hides plain or encoded donor storage. */
+export interface KvDonorAttention extends KvAttentionView {
+  readonly width: number;
+  readonly dtype: Dtype;
+  readonly offsets: readonly number[];
+  readonly starts: readonly number[];
+  readonly ends: readonly number[];
+}
 
 export interface Cache {
+  /** Maximum committed positions before this state changes precision. */
+  maxAppendTokens?(): number;
+  captureDonorRows?(): KvDonorRows;
+  captureDonorAttention?(): KvDonorAttention;
+  /** A preparation method commits each request's own precision boundary.
+   * Splitting a forward for a sibling must not move that boundary. */
+  readonly prefillMaintenance?: {
+    beginPrefill(): void;
+    commitPrefill(rows: readonly number[]): void;
+    endPrefill(): void;
+  };
+  /** Optional representation-owned affine conversion. Logical position can
+   * differ from the physical write head in a padded rotating layout. */
+  readonly affineConversion?: { readonly offset: number; toQuantized(groupSize: number, bits: number): Cache };
+  readonly turboConversion?: { readonly offset: number; toTurboQuantized(kBits: number, vBits: number): Cache };
+  readonly attentionState?: KvAttentionState;
+  /** Earliest prefix whose retained representation can resume generation.
+   * Irreversible transitions may preserve bytes while invalidating older
+   * precision boundaries. Reuse policy must not trim below this offset. */
+  minimumReusableOffset?: number;
+  readonly rotatedValueAttention?: RotatedValueAttentionState;
+  readonly quantizedAttention?: QuantizedAttentionState;
   offset: number;
   /** Stable storage identity for compatibility guards and persistence.
    *  REQUIRED: every capability predicate (isPlainKvCache, isRotating*,
@@ -364,6 +447,19 @@ export interface BatchableCache extends Cache {
   projectedBytes(tokens: number): number;
 }
 
+/** Padding is interpreted by cache geometry, independently of scheduling.
+ * Left padding initializes an empty group; lengths count input columns before
+ * trailing padding. Finalize before ordinary decode or publishing checkpoints. */
+export interface PrefillPadding {
+  readonly leftPadding?: readonly number[];
+  readonly lengths: readonly number[];
+  readonly rightPadding?: readonly number[];
+}
+export interface PaddedPrefillCache extends Cache {
+  preparePrefill(padding: PrefillPadding): void;
+  finalizePrefill(): void;
+}
+
 export function isBatchableCache(cache: Cache): cache is BatchableCache {
   const candidate = cache as Partial<BatchableCache>;
   return (
@@ -487,6 +583,13 @@ export class KVCache implements Cache {
   }
 
   /** 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];
@@ -578,6 +681,7 @@ export class KVCache implements Cache {
  *  quantization is NYI upstream; sliding layers are window-capped
  *  anyway). Attention dispatches to quantizedSdpa for these. */
 export class QuantizedKVCache implements Cache {
+  minimumReusableOffset?: number;
   static readonly STEP = 256;
   /** Set only by compiled-decode trace adapters (see Cache). */
   readonly ropeOffsetArr?: MlxArray;
@@ -587,6 +691,8 @@ export class QuantizedKVCache implements Cache {
 
   constructor(readonly groupSize: number, readonly bits: number) {}
 
+  get quantizedAttention(): QuantizedAttentionState { return this; }
+
   signature(): string { return `kv:quant:${this.bits}:${this.groupSize}`; }
 
   bytesPerToken(): number {
@@ -774,6 +880,16 @@ export class QuantizedKVCache implements Cache {
   /** Chronological (K, V) triples sliced to offset — the quantized twin of
    *  KVCache.temporalView (batched merge/extend/filter reads). Caller owns
    *  the returned views. */
+  captureDonorAttention(): KvDonorAttention {
+    const [keys, values] = this.temporalView();
+    const B = keys.packed.shape[0]!, width = keys.packed.shape[2]!;
+    return { width, dtype: keys.scales.dtype, offsets: Array(B).fill(this.offset),
+      starts: Array(B).fill(0), ends: Array(B).fill(width),
+      attend: (q, scale, mask) => quantizedSdpa(q, keys, values, scale, mask, this.groupSize, this.bits),
+      dispose() { disposeTriple(keys); disposeTriple(values); },
+    };
+  }
+
   temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
     if (!this.keys || !this.values) throw new Error("cache is empty");
     const cut = (t: ops.QuantizedTensor): ops.QuantizedTensor => {
@@ -798,6 +914,7 @@ export class QuantizedKVCache implements Cache {
  *  entries, so decode attends over at most the window. RoPE offsets use
  *  the true position; masks use the buffer-clamped offset. */
 export class RotatingKVCache implements Cache {
+  declare minimumReusableOffset?: number;
   static readonly STEP = 256;
   /** Set only by compiled-decode trace adapters (see Cache). */
   readonly ropeOffsetArr?: MlxArray;
@@ -1065,6 +1182,13 @@ export class RotatingKVCache implements Cache {
 
   /** Chronological (K, V) view, valid length min(offset, maxSize)
    *  (port of optiq kv_view._read_cache_temporal). */
+  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 tk = this.#temporalOrder(this.keys);
@@ -1133,6 +1257,8 @@ export const disposeTriple = (t: ops.QuantizedTensor): void => {
  *  + SDPA patches are unnecessary here: our SharedKv carries
  *  groupSize/bits through the donor→sharer plumbing explicitly. */
 export class RotatingQuantizedKVCache implements Cache {
+  get quantizedAttention(): QuantizedAttentionState { return this; }
+  declare minimumReusableOffset?: number;
   static readonly STEP = 256;
   /** Set only by compiled-decode trace adapters (see Cache). */
   readonly ropeOffsetArr?: MlxArray;
@@ -1390,6 +1516,16 @@ export class RotatingQuantizedKVCache implements Cache {
   /** Chronological (K, V) triples cut to the valid window — the quantized
    *  twin of RotatingKVCache.temporalView (batched merge reads, Phase 3
    *  milestone 2). Caller owns the returned views. */
+  captureDonorAttention(): KvDonorAttention {
+    const [keys, values] = this.temporalView();
+    const B = keys.packed.shape[0]!, width = keys.packed.shape[2]!;
+    return { width, dtype: keys.scales.dtype, offsets: Array(B).fill(this.offset),
+      starts: Array(B).fill(0), ends: Array(B).fill(width),
+      attend: (q, scale, mask) => quantizedSdpa(q, keys, values, scale, mask, this.groupSize, this.bits),
+      dispose() { disposeTriple(keys); disposeTriple(values); },
+    };
+  }
+
   temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
     if (!this.keys || !this.values) throw new Error("cache is empty");
     const valid = Math.min(this.offset, this.maxSize);
@@ -1504,28 +1640,6 @@ export class RotatingQuantizedKVCache implements Cache {
   }
 }
 
-/** One TurboQuant-encoded (K, V) storage tuple — the 5 arrays kv-store.ts
- *  and toQuantized/fromKVCache pass around. Not `ops.QuantizedTensor`
- *  (mlx's affine int4/int8 scheme): this is the rotation + Lloyd-Max
- *  layout (docs/design/turboquant.md), asymmetric-affine for keys,
- *  FWHT+Lloyd-Max for values. Field order is the kv-store.ts tensor-slot
- *  contract — do not reorder without updating snapshotCache/loadKvCache. */
-export interface TurboQuantTensor {
-  kIdx: MlxArray;
-  kScales: MlxArray;
-  kZeros: MlxArray;
-  vPacked: MlxArray;
-  vScales: MlxArray;
-}
-
-export const disposeTurboQuant = (t: TurboQuantTensor): void => {
-  t.kIdx.dispose();
-  t.kScales.dispose();
-  t.kZeros.dispose();
-  t.vPacked.dispose();
-  t.vScales.dispose();
-};
-
 /** TurboQuant KV cache — v1 (docs/design/turboquant.md): dequantize-
  *  on-fetch. Deliberately does NOT subclass KVCache/RotatingKVCache (or
  *  QuantizedKVCache) — a novel class fails every generated-file
@@ -1545,6 +1659,7 @@ export const disposeTurboQuant = (t: TurboQuantTensor): void => {
  *  a full-window dequant every step; the deferred-InvFWHT trick is a
  *  documented non-goal until the quality gate passes. */
 export class TurboQuantKVCache implements Cache {
+  minimumReusableOffset = 0;
   readonly stateNeedsDispose = true;
   static readonly STEP = 256;
   /** Set only by compiled-decode trace adapters (see Cache) — TurboQuant
@@ -1557,10 +1672,23 @@ export class TurboQuantKVCache implements Cache {
   /** Validated lazily on first update, once head_dim is known. */
   #headDim: number | null = null;
   /** Experimental operation selection, captured once for this cache. */
-  readonly #fusedDecode = process.env.MLX_BUN_TURBOQUANT_FUSED_DECODE === "1";
+  readonly #codec: KvCodec<TurboQuantTensor>;
 
-  constructor(readonly kBits: number, readonly vBits: number) {}
+  constructor(readonly kBits: number, readonly vBits: number,
+    readonly fusedDecode = turboQuantFusedDecode()) {
+    this.#codec = new TurboQuantCodec(kBits, vBits, fusedDecode);
+  }
+
+  captureDonorRows(): KvDonorRows {
+    if (!this.#kv || this.#headDim === null) throw new Error("cache is empty");
+    const [keys, values] = this.#decode(this.#kv, this.offset, this.#headDim, false);
+    const B = keys.shape[0]!;
+    return { keys, values, offsets: Array(B).fill(this.offset),
+      starts: Array(B).fill(0), ends: Array(B).fill(this.offset) };
+  }
+  captureDonorAttention(): KvDonorAttention { return decodedKvDonorAttention(this.captureDonorRows()); }
 
+  get rotatedValueAttention(): RotatedValueAttentionState { return this; }
   signature(): string { return `kv:turboquant:${this.kBits}:${this.vBits}`; }
 
   bytesPerToken(): number {
@@ -1612,55 +1740,9 @@ export class TurboQuantKVCache implements Cache {
     return this.#kv ? this.#kv.kIdx.shape[2]! : 0;
   }
 
-  /** Encode k/v [B,H,L,D] into a fresh TurboQuantTensor (packed at rest). */
-  #encode(k: MlxArray, v: MlxArray): TurboQuantTensor {
-    const signed = this.kBits === 8;
-    const kEnc = tq.encodeKeys(k, this.kBits, signed);
-    const kIdxPacked = this.kBits < 8 ? tq.packBits(kEnc.indices, this.kBits) : kEnc.indices;
-    if (kIdxPacked !== kEnc.indices) kEnc.indices.dispose();
-    const vEnc = tq.encodeValues(v, this.vBits);
-    const vIdxPacked = this.vBits < 8 ? tq.packBits(vEnc.indices, this.vBits) : vEnc.indices;
-    if (vIdxPacked !== vEnc.indices) vEnc.indices.dispose();
-    return {
-      kIdx: kIdxPacked, kScales: kEnc.scales, kZeros: kEnc.zeros,
-      vPacked: vIdxPacked, vScales: vEnc.scales,
-    };
-  }
-
-  /** Decode a TurboQuantTensor's active [B,H,upTo,*] window back to
-   *  bf16 [B,H,upTo,headDim] (k, v). Caller owns the returned arrays.
-   *  deferV leaves V in the rotated domain (decodeValuesRotated) — the
-   *  caller un-rotates its attention output instead (tq.unrotateValues). */
+  #encode(k: MlxArray, v: MlxArray): TurboQuantTensor { return this.#codec.encode(k, v); }
   #decode(t: TurboQuantTensor, upTo: number, headDim: number, deferV = false): [MlxArray, MlxArray] {
-    const cut = (a: MlxArray): MlxArray => {
-      const [B, H, , D] = a.shape as [number, number, number, number];
-      return a.slice([0, 0, 0, 0], [B, H, upTo, D]);
-    };
-    if (this.#fusedDecode) {
-      const inputs = [cut(t.kIdx), cut(t.kScales), cut(t.kZeros), cut(t.vPacked), cut(t.vScales)] as const;
-      try {
-        const decoded = tq.tryDecodePackedKv(inputs, this.kBits, this.vBits, headDim, deferV);
-        if (decoded) return decoded;
-      } finally { for (const input of inputs) input.dispose(); }
-    }
-    const kIdxCut = cut(t.kIdx);
-    const kIdxUnpacked = this.kBits < 8 ? tq.unpackBits(kIdxCut, this.kBits, headDim) : kIdxCut;
-    if (kIdxUnpacked !== kIdxCut) kIdxCut.dispose();
-    const kScalesCut = cut(t.kScales);
-    const kZerosCut = cut(t.kZeros);
-    const k = tq.decodeKeys(kIdxUnpacked, kScalesCut, kZerosCut);
-    for (const a of [kIdxUnpacked, kScalesCut, kZerosCut]) a.dispose();
-
-    const vPackedCut = cut(t.vPacked);
-    const vIdxUnpacked = this.vBits < 8 ? tq.unpackBits(vPackedCut, this.vBits, headDim) : vPackedCut;
-    if (vIdxUnpacked !== vPackedCut) vPackedCut.dispose();
-    const vScalesCut = cut(t.vScales);
-    const v = deferV
-      ? tq.decodeValuesRotated(vIdxUnpacked, vScalesCut, this.vBits)
-      : tq.decodeValues(vIdxUnpacked, vScalesCut, this.vBits);
-    for (const a of [vIdxUnpacked, vScalesCut]) a.dispose();
-
-    return [k, v];
+    return this.#codec.decode(t, upTo, headDim, deferV);
   }
 
   /** Quantize the newly-appended k/v, append into growable packed
@@ -1805,6 +1887,7 @@ export class TurboQuantKVCache implements Cache {
       q.#headDim = D;
     }
     q.offset = cache.offset;
+    q.minimumReusableOffset = (cache as Cache).minimumReusableOffset ?? 0;
     cache.dispose();
     return q;
   }
@@ -1813,6 +1896,7 @@ export class TurboQuantKVCache implements Cache {
     if (this.#kv) disposeTurboQuant(this.#kv);
     this.#kv = null;
     this.offset = 0;
+    this.minimumReusableOffset = 0;
   }
 }
 
@@ -1860,8 +1944,23 @@ export function quantizedSdpaUnfused(
     vT = expand(vq);
   }
 
-  let scores = ops.quantizedMatmulQT(queries, kT, true, groupSize, bits);
+  // M4 Pro's native key matvec retains exact arithmetic when three GQA
+  // heads share a batch. Restore score geometry before softmax/value work.
+  const groupHeads = B <= 2 && H === 24 && KV === 4 && L === 3 && D === 256 &&
+    N >= 8192 && groupSize === 64 && bits === 4 &&
+    (q.dtype === Dtype.bfloat16 || q.dtype === Dtype.float32) &&
+    deviceArchitecture() === "applegpu_g16s";
+  let keyQueries = queries;
+  if (groupHeads) {
+    keyQueries = ops.reshape(queries, [B, KV, 2, 9, D]);
+    owned.push(keyQueries);
+  }
+  let scores = ops.quantizedMatmulQT(keyQueries, kT, true, groupSize, bits);
   owned.push(scores);
+  if (groupHeads) {
+    scores = ops.reshape(scores, [B, KV, nRep, L, N]);
+    owned.push(scores);
+  }
 
   let maskArr: MlxArray | null = null;
   let ownsMask = false;
@@ -1885,9 +1984,11 @@ export function quantizedSdpaUnfused(
     }
   }
   if (maskArr) {
-    const ninf = ops.scalarLike(FINFO_MIN[scores.dtype] ?? -3.4e38, scores);
-    const masked = ops.where(maskArr, scores, ninf);
-    ninf.dispose();
+    let masked: MlxArray;
+    if (maskArr.dtype === Dtype.bool) {
+      using ninf = ops.scalarLike(FINFO_MIN[scores.dtype] ?? -3.4e38, scores);
+      masked = ops.where(maskArr, scores, ninf);
+    } else masked = ops.add(scores, maskArr);
     if (ownsMask) maskArr.dispose();
     owned.push(masked);
     scores = masked;

package/src/model/gemma4.ts

@@ -45,18 +45,19 @@ import {
   RMSNorm,
   RotatingKVCache,
   RotatingQuantizedKVCache,
-  TurboQuantKVCache,
   quantizedSdpa,
   type Cache,
   type Mask,
   type SharedKv,
+  captureRopeOffsets,
   type LoraWeights,
 } from "./gemma4-base";
+import { mapPackedTokens, type TokenGroup } from "./token-groups";
 import { Checkpoint } from "../mlx/checkpoint";
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 import { unrotateValues as tqUnrotateValues } from "../mlx/turboquant-ops";
 import { CompiledFunction } from "../mlx/compile";
-import { flagOn } from "../runtime-config";
+import { runtimeFlag } from "../runtime-config";
 
 /** Verbatim port of mlx_lm/models/gemma4_text.py:
  *    `@partial(mx.compile, shapeless=True) def geglu(gate, x): return nn.gelu_approx(gate) * x`
@@ -77,7 +78,7 @@ import { flagOn } from "../runtime-config";
  *  ~9-op path. `MLX_BUN_COMPILED_GEGLU=0` (--compiled-activations off) selects the
  *  uncompiled composition — same L1 parity, slower (the A/B opt-out). */
 export function compiledGegluActive(): boolean {
-  return flagOn("MLX_BUN_COMPILED_GEGLU", true);
+  return runtimeFlag("MLX_BUN_COMPILED_GEGLU", true);
 }
 
 let _gegluClosure: CompiledFunction | null = null;
@@ -304,15 +305,11 @@ class Attention {
       }
 
       const offset = cache.offset;
-      // Capture the RoPE offset array ONCE (like `offset` above): the
-      // updateAndFetch below advances cache.offset, so re-reading
-      // cache.ropeOffsetArr for Q *after* the write would hand K and Q
-      // different positions. Harmless today (real caches leave it unset;
-      // compiled-decode passes a constant trace input), but a hard
-      // prerequisite for batched decode, where ropeOffsetArr carries per-row
-      // positions derived from the pre-write offset and MUST be identical for
-      // this step's K and Q.
-      const offsetArr = cache.ropeOffsetArr;
+      // Own the pre-write positions through every query and shared-KV
+      // consumer. Advancing a row cache may replace and release its borrowed
+      // offset array before Q's RoPE is built. The fetched shared state owns
+      // this retained handle and releases it with K/V at the end of the pass.
+      const offsetArr = captureRopeOffsets(cache);
 
       const kNormed = this.kNorm!.forward(k);
       const kT = ops.transposeAxes(kNormed, [0, 2, 1, 3]);
@@ -328,23 +325,29 @@ class Attention {
       if (v !== k) v.dispose();
       k.dispose();
 
-      if (cache instanceof QuantizedKVCache || cache instanceof RotatingQuantizedKVCache) {
-        const [kq, vq] = cache.updateAndFetchQuantized(kRoped, vT);
+      if (cache.attentionState) {
+        const attention = cache.attentionState.appendAndFetch(kRoped, vT);
+        kRoped.dispose(); vT.dispose();
+        shared = { kind: "view", attention, offset, offsetArr };
+      } else if (cache.quantizedAttention) {
+        const quantized = cache.quantizedAttention;
+        const [kq, vq] = quantized.updateAndFetchQuantized(kRoped, vT);
         kRoped.dispose();
         vT.dispose();
         shared = {
           kind: "quant", keys: kq, values: vq, offset,
-          groupSize: cache.groupSize, bits: cache.bits,
+          groupSize: quantized.groupSize, bits: quantized.bits,
           offsetArr,
         };
-      } else if (cache instanceof TurboQuantKVCache) {
+      } else if (cache.rotatedValueAttention) {
         // Deferred-V read: values stay rotated; the attention output is
         // un-rotated below (shared.vRotated) — including by KV-shared
         // consumer layers, which see the flag through sharedIn.
-        const [keys, values] = cache.updateAndFetchDeferredV(kRoped, vT);
+        const [keys, values] = cache.rotatedValueAttention.updateAndFetchDeferredV(kRoped, vT);
         kRoped.dispose();
         vT.dispose();
-        shared = { kind: "plain", keys, values, offset, offsetArr, vRotated: true };
+        shared = { kind: "plain", keys, values, offset, offsetArr, vRotated: true,
+          restoreValues: cache.rotatedValueAttention.captureValueTransform?.() };
       } else {
         const [keys, values] = cache.updateAndFetch(kRoped, vT);
         kRoped.dispose();
@@ -360,7 +363,9 @@ class Attention {
 
     let attn: MlxArray;
     const ta = getTrainingAttn();
-    if (shared.kind === "quant") {
+    if (shared.kind === "view") {
+      attn = shared.attention.attend(q, 1.0, mask);
+    } else if (shared.kind === "quant") {
       attn = quantizedSdpa(q, shared.keys, shared.values, 1.0, mask, shared.groupSize, shared.bits);
     } else if (
       ta === "flash" && flashSupported(q) && (mask.mode === "causal" || mask.mode === "array")
@@ -383,7 +388,7 @@ class Attention {
     if (shared.kind === "plain" && shared.vRotated) {
       // Attention is linear in V: one InvFWHT on the [B,H,L,D] output
       // replaces one per cached token (TurboQuant deferred-V).
-      attn = disposing(attn, tqUnrotateValues(attn));
+      attn = disposing(attn, shared.restoreValues ? shared.restoreValues(attn) : tqUnrotateValues(attn));
     }
     const attnT = ops.transposeAxes(attn, [0, 2, 1, 3]);
     attn.dispose();
@@ -851,6 +856,13 @@ export class Gemma4Model {
     );
   }
 
+  /** Computed model constants must retain their constructor arithmetic when
+   * a restored request enters compiled decode without a prefill first. */
+  materializeGraphConstants(): void {
+    const constants = this.layers.flatMap(layer => layer.attn.ropeFreqs ? [layer.attn.ropeFreqs] : []);
+    if (constants.length) ops.evalAll(constants);
+  }
+
   /** ids [1, L] → final-norm hidden states [1, L, hidden]. */
   forwardHidden(ids: MlxArray, cache: Cache[]): MlxArray {
     let h = this.embed.encode(ids);
@@ -1000,13 +1012,7 @@ export class Gemma4Model {
     for (let i = 0; i < this.numDonors; i++) {
       const s = intermediates[i];
       if (!s) continue;
-      if (s.kind === "plain") {
-        s.keys.dispose();
-        s.values.dispose();
-      } else {
-        for (const t of [s.keys, s.values])
-          for (const a of [t.packed, t.scales, t.biases]) a.dispose();
-      }
+      Gemma4Model.disposeSharedKv(s);
     }
     for (const m of masks.values()) m.arr?.dispose();
 
@@ -1015,6 +1021,85 @@ export class Gemma4Model {
     return h;
   }
 
+  /** Mixed token execution: attention retains each group's cache, positions and
+   * SDPA shape; feed-forward work packs only real tokens. Kernel selection for
+   * packed matmuls can differ from a solo GEMV, so this is an explicit Lab port
+   * until its matching packed oracle and serving comparison are accepted. */
+  forwardHiddenMixed(work: readonly TokenGroup[]): MlxArray[] {
+    if (work.length === 1 && !work[0]!.captureLayer) return [this.forwardHidden(work[0]!.ids, work[0]!.cache)];
+    const groups: Array<{ ids: MlxArray; cache: Cache[]; h: MlxArray;
+      masks: Map<string, Mask>; perLayer: MlxArray | null; intermediates: (SharedKv | null)[] }> = [];
+    const results: MlxArray[] = [];
+    const packedMlp = runtimeFlag("MLX_BUN_MIXED_PACKED_MLP", true) && !work.some(group => group.preserveTokenGeometry);
+    try {
+      for (const { ids, cache } of work) {
+        const group = { ids, cache, h: this.embed.encode(ids), masks: new Map<string, Mask>(),
+          perLayer: null as MlxArray | null, intermediates: Array<SharedKv | null>(this.layers.length).fill(null) };
+        groups.push(group);
+        group.h = disposing(group.h, ops.mulScalar(group.h, this.embedScale));
+        for (let i = 0; i < this.numDonors; i++) {
+          const type = this.layers[i]!.layerType;
+          if (!group.masks.has(type)) group.masks.set(type,
+            cache[i]!.makeMask(ids.shape[1]!, type === "sliding_attention" ? this.windowSize : null));
+        }
+        if (this.perLayerWidth > 0) group.perLayer = this.computePerLayerInputs(ids, group.h);
+      }
+      for (let i = 0; i < this.layers.length; i++) {
+        const layer = this.layers[i]!, ci = this.cacheIndex[i]!;
+        const mids: MlxArray[] = [], inputs: MlxArray[] = [];
+        try {
+          for (const group of groups) {
+            const sharedIn = ci === -1 ? group.intermediates[this.previousKvs[i]!]! : null;
+            const { h, shared } = layer.forwardAttn(group.h, group.masks.get(layer.layerType)!,
+              ci === -1 ? null : group.cache[ci]!, sharedIn);
+            mids.push(h); group.intermediates[i] = shared;
+            if (group.perLayer) {
+              const [B, L] = group.ids.shape;
+              using part = group.perLayer.slice([0, 0, i, 0], [B!, L!, i + 1, this.perLayerWidth]);
+              inputs.push(ops.reshape(part, [B!, L!, this.perLayerWidth]));
+            }
+          }
+          // Per-layer inputs use the same token order as the hidden states.
+          let packedInput: MlxArray | null = null;
+          if (packedMlp && inputs.length) {
+            const flat = inputs.map(input => ops.reshape(input, [1, -1, this.perLayerWidth]));
+            try { packedInput = ops.concatAxis(flat, 1); }
+            finally { for (const input of flat) input.dispose(); }
+          }
+          using pli = packedInput;
+          const outputs: MlxArray[] = [];
+          try {
+            if (packedMlp) outputs.push(...mapPackedTokens(mids, packed => layer.forwardMlp(packed, pli)));
+            else for (const [row, mid] of mids.entries()) outputs.push(layer.forwardMlp(mid, inputs[row] ?? null));
+          } catch (error) { for (const output of outputs) output.dispose(); throw error; }
+          for (const [index, group] of groups.entries()) {
+            group.h.dispose(); group.h = outputs[index]!;
+          }
+          for (const [index, group] of groups.entries()) work[index]!.captureLayer?.(i, group.h);
+        } finally {
+          for (const mid of mids) mid.dispose();
+          for (const input of inputs) input.dispose();
+        }
+      }
+      for (const [index, group] of groups.entries()) {
+        const hidden = this.finalNorm.forward(group.h); results.push(hidden);
+        work[index]!.captureLayer?.(this.layers.length, hidden);
+      }
+      return results;
+    } catch (error) { for (const result of results) result.dispose(); throw error; }
+    finally {
+      for (const group of groups) {
+        group.h.dispose(); group.perLayer?.dispose();
+        for (let i = 0; i < this.numDonors; i++) {
+          const shared = group.intermediates[i];
+          if (shared) Gemma4Model.disposeSharedKv(shared);
+        }
+        for (const mask of group.masks.values()) mask.arr?.dispose();
+      }
+    }
+  }
+
+
   // --- Segmented-backward support (docs/design/orpo-training.md
   // §4 Phase B). These are ADDITIVE — forwardLayers is untouched. The segmented
   // driver builds masks + per-layer inputs ONCE, then drives runLayerRange per
@@ -1056,7 +1141,10 @@ export class Gemma4Model {
   /** Dispose a fetched donor K/V (plain or quantized) — mirrors forwardLayers'
    *  donor cleanup. */
   private static disposeSharedKv(s: SharedKv): void {
-    if (s.kind === "plain") {
+    s.offsetArr?.dispose();
+    if (s.kind === "view") {
+      s.attention.dispose();
+    } else if (s.kind === "plain") {
       s.keys.dispose();
       s.values.dispose();
     } else {

package/src/model/generated/gemma4-12b.ts

@@ -34,7 +34,10 @@ export const FINGERPRINT = "9f812d2eb461fcbe";
 export let generatedForwardUses = 0;
 
 function disposeSharedKv(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {

package/src/model/generated/gemma4-26b.ts

@@ -34,7 +34,10 @@ export const FINGERPRINT = "c9dd67ed5a525231";
 export let generatedForwardUses = 0;
 
 function disposeSharedKv(s: SharedKv): void {
-  if (s.kind === "plain") {
+  if (s.kind === "view") {
+    s.attention.dispose();
+    s.offsetArr?.dispose();
+  } else if (s.kind === "plain") {
     s.keys.dispose();
     s.values.dispose();
   } else {