mlx-bun 0.0.10 → 0.0.11

package/src/model/gemma4-base.ts

@@ -14,9 +14,7 @@ import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 import { expertOffloadArray } from "../expert-offload";
-import {
-  fusedDecodeKernelSupported, fusedDecodeSdpa, perfKernelEnabled,
-} from "./fused-decode-kernel";
+import * as tq from "../mlx/turboquant-ops";
 
 export type MaskMode = "" | "causal";
 export interface Mask {
@@ -206,7 +204,13 @@ export class QuantizedEmbedding {
  *  position (see src/model/compiled-decode.ts). */
 export type SharedKv =
   | { kind: "plain"; keys: MlxArray; values: MlxArray; offset: number;
-      offsetArr?: MlxArray }
+      offsetArr?: MlxArray;
+      /** TurboQuant deferred-V: `values` are still in the rotated (FWHT)
+       *  domain — every consumer must un-rotate its attention OUTPUT
+       *  (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 }
   | { kind: "quant"; keys: ops.QuantizedTensor; values: ops.QuantizedTensor;
       offset: number; groupSize: number; bits: number; offsetArr?: MlxArray };
 
@@ -306,9 +310,16 @@ export class KVCache implements Cache {
 
   makeMask(N: number, windowSize: number | null): Mask {
     if (N === 1) return { mode: "", arr: null };
-    if (this.offset === 0 && windowSize === null) return { mode: "causal", arr: null };
-    if (this.offset === 0 && windowSize !== null && N <= windowSize)
-      return { mode: "causal", 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) };
   }
 
@@ -508,14 +519,13 @@ export class QuantizedKVCache implements Cache {
 
   makeMask(N: number, windowSize: number | null): Mask {
     if (N === 1) return { mode: "", arr: null };
-    if (this.offset === 0 && (windowSize === null || N <= windowSize))
+    // Same rule as KVCache.makeMask (mlx-lm cache.py:114-125): windowless
+    // multi-token chunks are the string "causal" at any offset. The old
+    // causalEquivalent escape hatch existed precisely because this case
+    // used to ship an array; only windowed continuations materialize now.
+    if (windowSize === null || (this.offset === 0 && N <= windowSize))
       return { mode: "causal", arr: null };
-    return {
-      mode: "array",
-      arr: createCausalMask(N, this.offset, windowSize),
-      // a windowless continuation mask is exactly mlx-lm's "causal"
-      causalEquivalent: windowSize === null,
-    };
+    return { mode: "array", arr: createCausalMask(N, this.offset, windowSize) };
   }
 
   /** Compiled decode: growth half of updateAndFetchQuantized (L=1),
@@ -599,6 +609,21 @@ export class QuantizedKVCache implements Cache {
     this.offset = offset;
   }
 
+  /** Chronological (K, V) triples sliced to offset — the quantized twin of
+   *  KVCache.temporalView (batched merge/extend/filter reads). Caller owns
+   *  the returned views. */
+  temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    if (!this.keys || !this.values) throw new Error("cache is empty");
+    const cut = (t: ops.QuantizedTensor): ops.QuantizedTensor => {
+      const f = (a: MlxArray): MlxArray => {
+        const [b, h, , d] = a.shape as [number, number, number, number];
+        return a.slice([0, 0, 0, 0], [b, h, this.offset, d]);
+      };
+      return { packed: f(t.packed), scales: f(t.scales), biases: f(t.biases) };
+    };
+    return [cut(this.keys), cut(this.values)];
+  }
+
   dispose(): void {
     for (const a of this.state()) a.dispose();
     this.keys = this.values = null;
@@ -1181,6 +1206,25 @@ export class RotatingQuantizedKVCache implements Cache {
     this.ringIdx = idx;
   }
 
+  /** 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. */
+  temporalView(): [ops.QuantizedTensor, ops.QuantizedTensor] {
+    if (!this.keys || !this.values) throw new Error("cache is empty");
+    const valid = Math.min(this.offset, this.maxSize);
+    const cutValid = (t: ops.QuantizedTensor): ops.QuantizedTensor =>
+      mapTriple(t, (a) => {
+        const [B, H, , D] = a.shape as [number, number, number, number];
+        return a.slice([0, 0, 0, 0], [B, H, valid, D]);
+      });
+    const tk = this.#temporalOrder(this.keys);
+    const tv = this.#temporalOrder(this.values);
+    const out: [ops.QuantizedTensor, ops.QuantizedTensor] = [cutValid(tk), cutValid(tv)];
+    disposeTriple(tk);
+    disposeTriple(tv);
+    return out;
+  }
+
   /** Oracle: to_quantized on an already-quantized rotating cache is
    *  idempotent. */
   toQuantized(): RotatingQuantizedKVCache {
@@ -1279,6 +1323,297 @@ 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-kv.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-kv.md): dequantize-
+ *  on-fetch. Deliberately does NOT subclass KVCache/RotatingKVCache (or
+ *  QuantizedKVCache) — a novel class fails every generated-file
+ *  `#matches()` guard and every batching `instanceof` allow-list, so it
+ *  gets solo-only routing and monolith fallback for free, the same way
+ *  QuantizedKVCache does today. Only full-attention (non-rotating)
+ *  layers use this in v1; sliding-window layers stay bf16.
+ *
+ *  Storage: 5 arrays per the vllm-metal reference's `turbo_quant_encode`
+ *  tuple — kIdx (int8 unpacked at kBits=8, else bit-packed uint8),
+ *  kScales/kZeros (fp16, one pair per 32-group), vPacked (bit-packed
+ *  uint8, no-op at vBits=8), vScales (fp16, one per 32-group).
+ *  updateAndFetch quantizes ONLY the newly appended tokens (unlike
+ *  QuantizedKVCache's whole-buffer-toQuantized conversion path — the
+ *  streaming append is the whole point of a live cache) and returns the
+ *  DEQUANTIZED bf16 active window so ops.sdpa runs unmodified — v1 pays
+ *  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 {
+  static readonly STEP = 256;
+  /** Set only by compiled-decode trace adapters (see Cache) — TurboQuant
+   *  is not a compiled-decode participant in v1 (novel class, monolith
+   *  fallback only), so this is always unset; captured once per step by
+   *  construction, never refreshed mid-step (batched-rotating-quant.ts:55-63). */
+  readonly ropeOffsetArr?: MlxArray;
+  #kv: TurboQuantTensor | null = null;
+  offset = 0;
+  /** Validated lazily on first update, once head_dim is known. */
+  #headDim: number | null = null;
+
+  constructor(readonly kBits: number, readonly vBits: number) {}
+
+  get headDim(): number | null {
+    return this.#headDim;
+  }
+
+  #validateHeadDim(dim: number): void {
+    if (this.#headDim !== null) return;
+    const supported = [64, 128, 256, 512];
+    if (dim % 32 !== 0 || !supported.includes(dim)) {
+      throw new Error(
+        `TurboQuantKVCache: head_dim ${dim} must be divisible by 32 and in {${supported.join(",")}}`,
+      );
+    }
+    this.#headDim = dim;
+  }
+
+  #packedKDim(dim: number): number {
+    return tq.packedDim(dim, this.kBits);
+  }
+  #packedVDim(dim: number): number {
+    return tq.packedDim(dim, this.vBits);
+  }
+  #nGroups(dim: number): number {
+    return dim / tq.TURBOQUANT_BLOCK_SIZE;
+  }
+
+  #alloc(B: number, H: number, T: number, dim: number): TurboQuantTensor {
+    const kIdxDtype = this.kBits === 8 ? Dtype.int8 : Dtype.uint8;
+    return {
+      kIdx: ops.zeros([B, H, T, this.#packedKDim(dim)], kIdxDtype),
+      kScales: ops.zeros([B, H, T, this.#nGroups(dim)], Dtype.float16),
+      kZeros: ops.zeros([B, H, T, this.#nGroups(dim)], Dtype.float16),
+      vPacked: ops.zeros([B, H, T, this.#packedVDim(dim)], Dtype.uint8),
+      vScales: ops.zeros([B, H, T, this.#nGroups(dim)], Dtype.float16),
+    };
+  }
+
+  #seqLen(): number {
+    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). */
+  #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]);
+    };
+    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];
+  }
+
+  /** Quantize the newly-appended k/v, append into growable packed
+   *  storage (256-step growth, mirroring KVCache), advance offset, then
+   *  dequantize the WHOLE active window (v1 semantics) so the standard
+   *  ops.sdpa path runs unmodified. */
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const D = this.#append(k, v);
+    return this.#decode(this.#kv!, this.offset, D);
+  }
+
+  /** Deferred-inverse-FWHT read path: same append, but the returned V
+   *  window stays in the ROTATED domain. The caller MUST un-rotate its
+   *  attention output with tq.unrotateValues (see SharedKv.vRotated) —
+   *  same decode by linearity, O(q·d log d) per step instead of
+   *  O(T·d log d). Opt-in per attention site; sites that don't know about
+   *  rotation keep calling updateAndFetch and stay correct. */
+  updateAndFetchDeferredV(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const D = this.#append(k, v);
+    return this.#decode(this.#kv!, this.offset, D, true);
+  }
+
+  /** Shared append: grow storage, encode the new tokens, slice-write them
+   *  at the previous offset, advance offset. Returns head_dim. */
+  #append(k: MlxArray, v: MlxArray): number {
+    const [B, H, L, D] = k.shape as [number, number, number, number];
+    this.#validateHeadDim(D);
+    const prev = this.offset;
+
+    if (!this.#kv || prev + L > this.#seqLen()) {
+      const nSteps = Math.floor((TurboQuantKVCache.STEP + L - 1) / TurboQuantKVCache.STEP);
+      const newSteps = nSteps * TurboQuantKVCache.STEP;
+      if (this.#kv && prev % TurboQuantKVCache.STEP !== 0) {
+        const trimTo = (t: TurboQuantTensor): TurboQuantTensor => {
+          const cut = (a: MlxArray): MlxArray => {
+            const [b, h, , d] = a.shape as [number, number, number, number];
+            const s = a.slice([0, 0, 0, 0], [b, h, prev, d]);
+            a.dispose();
+            return s;
+          };
+          return { kIdx: cut(t.kIdx), kScales: cut(t.kScales), kZeros: cut(t.kZeros),
+            vPacked: cut(t.vPacked), vScales: cut(t.vScales) };
+        };
+        this.#kv = trimTo(this.#kv);
+      }
+      const grown = this.#alloc(B, H, newSteps, D);
+      if (this.#kv) {
+        const cat = (a: MlxArray, b: MlxArray): MlxArray => {
+          const out = ops.concatAxis([a, b], 2);
+          a.dispose();
+          b.dispose();
+          return out;
+        };
+        this.#kv = {
+          kIdx: cat(this.#kv.kIdx, grown.kIdx),
+          kScales: cat(this.#kv.kScales, grown.kScales),
+          kZeros: cat(this.#kv.kZeros, grown.kZeros),
+          vPacked: cat(this.#kv.vPacked, grown.vPacked),
+          vScales: cat(this.#kv.vScales, grown.vScales),
+        };
+      } else {
+        this.#kv = grown;
+      }
+    }
+
+    this.offset += L;
+    const enc = this.#encode(k, v);
+    const write = (dst: MlxArray, src: MlxArray): MlxArray => {
+      const [b, h, , d] = dst.shape as [number, number, number, number];
+      const u = ops.sliceUpdate(dst, src, [0, 0, prev, 0], [b, h, this.offset, d]);
+      dst.dispose();
+      src.dispose();
+      return u;
+    };
+    this.#kv = {
+      kIdx: write(this.#kv!.kIdx, enc.kIdx),
+      kScales: write(this.#kv!.kScales, enc.kScales),
+      kZeros: write(this.#kv!.kZeros, enc.kZeros),
+      vPacked: write(this.#kv!.vPacked, enc.vPacked),
+      vScales: write(this.#kv!.vScales, enc.vScales),
+    };
+
+    return D;
+  }
+
+  /** Same rule as KVCache/QuantizedKVCache.makeMask (mlx-lm cache.py:114-125):
+   *  windowless multi-token chunks are the string "causal" at any offset —
+   *  TurboQuant v1 is full-attention-only, never windowed. */
+  makeMask(N: number, windowSize: number | null): Mask {
+    if (N === 1) return { mode: "", arr: null };
+    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) };
+  }
+
+  /** Trimmed-to-offset arrays in documented order: kIdx, kScales, kZeros,
+   *  vPacked, vScales (kv-store.ts tensor-slot ordering for kind
+   *  "turboquant"). Empty cache → []. */
+  state(): MlxArray[] {
+    if (!this.#kv) return [];
+    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, this.offset, D]);
+    };
+    return [cut(this.#kv.kIdx), cut(this.#kv.kScales), cut(this.#kv.kZeros),
+      cut(this.#kv.vPacked), cut(this.#kv.vScales)];
+  }
+
+  isTrimmable(): boolean {
+    return true;
+  }
+
+  /** Drop the last n tokens (byte-safe: groups run along head_dim, so a
+   *  token-axis slice never splits a group). Same front/back convention
+   *  as KVCache.trim: offset shrinks, storage is overwritten on next write. */
+  trim(n: number): void {
+    this.offset = Math.max(0, this.offset - n);
+  }
+
+  /** Adopt persisted/converted state (takes ownership of the arrays). */
+  restoreState(t: TurboQuantTensor, offset: number, headDim: number): void {
+    this.dispose();
+    this.#kv = t;
+    this.offset = offset;
+    this.#headDim = headDim;
+  }
+
+  /** Convert an existing bf16 KVCache/RotatingKVCache's live window in one
+   *  shot: quantize the whole [.., :offset, :] region, preserve offset,
+   *  dispose the source's arrays. Mirrors KVCache.toQuantized/
+   *  RotatingKVCache.toQuantized's contract (source cache is consumed). */
+  static fromKVCache(cache: KVCache | RotatingKVCache, kBits: number, vBits: number): TurboQuantKVCache {
+    const q = new TurboQuantKVCache(kBits, vBits);
+    const [k, v] = cache.keys && cache.values ? cache.temporalView() : [null, null];
+    if (k && v) {
+      const D = k.shape[3]!;
+      q.#validateHeadDim(D);
+      const enc = q.#encode(k, v);
+      k.dispose();
+      v.dispose();
+      q.#kv = enc;
+      q.#headDim = D;
+    }
+    q.offset = cache.offset;
+    cache.dispose();
+    return q;
+  }
+
+  dispose(): void {
+    if (this.#kv) disposeTurboQuant(this.#kv);
+    this.#kv = null;
+    this.offset = 0;
+  }
+}
+
 export const FINFO_MIN: Partial<Record<Dtype, number>> = {
   [Dtype.bfloat16]: -3.3895313892515355e38,
   [Dtype.float16]: -65504,
@@ -1338,6 +1673,14 @@ export function quantizedSdpaUnfused(
     for (const a of [qIdx, kIdx, qCol, kRow]) a.dispose();
   } else if (mask.mode === "array") {
     maskArr = mask.arr;
+    // Batched padding masks are [B,1,L,S] (buildBatchedDecodeMask). When GQA
+    // reshaped the scores to 5-D [B,KV,nRep,L,S], trailing-dim broadcasting
+    // would misalign the mask's B with KV — insert the unit axis explicitly:
+    // [B,1,1,L,S] broadcasts correctly over (KV, nRep).
+    if (maskArr && maskArr.ndim === 4 && nRep > 1) {
+      maskArr = ops.expandDims(maskArr, 1);
+      ownsMask = true;
+    }
   }
   if (maskArr) {
     const ninf = ops.scalarLike(FINFO_MIN[scores.dtype] ?? -3.4e38, scores);
@@ -1568,38 +1911,19 @@ function fusedSdpaSupported(q: MlxArray, mask: Mask, groupSize: number, bits: nu
 /** Quantized-cache SDPA dispatch: L > 1 (prefill/continuation) with a
  *  supported config goes through the N-tiled fused path; decode (L = 1)
  *  and unsupported configs stay on the stock unfused port. Exported for
- *  the dispatch-gate tests.
- *
- *  MLX_BUN_FUSED_DECODE=1 (NEXT UP 1b experiment): tile decode too,
- *  matching optiq's wrapper, which has no L gate — its serving decode
- *  over quantized caches runs N tiles per step, and Phase 15 measured
- *  its kv-mixed decode tax @8k as ~free where ours is ~3%. Off by
- *  default until a cleared-machine A/B shows it pays; directional
- *  numbers via scripts/bench-fused-decode.ts. */
+ *  the dispatch-gate tests. */
 export function quantizedSdpa(
   q: MlxArray, kq: ops.QuantizedTensor, vq: ops.QuantizedTensor,
   scale: number, mask: Mask, groupSize: number, bits: number,
 ): MlxArray {
-  // Perf mode (Phase E): the fused decode kernel takes supported L=1
-  // dispatches. NOT bit-exact (online softmax) — gated against the
-  // frozen perf oracle; compat (MLX_BUN_PERF_KERNEL=0) stays the -O0
-  // reference. Never inside a compiled trace: CustomKernel has no
-  // output_shapes, so it cannot live in a (shapeless) closure.
-  if (
-    mask.mode === "" && scale === 1.0 && perfKernelEnabled() &&
-    !isCompiledTrace() && fusedDecodeKernelSupported(q, bits, groupSize)
-  )
-    return fusedDecodeSdpa(q, kq, vq, groupSize, bits);
-  const tile = q.shape[2]! > 1 || process.env.MLX_BUN_FUSED_DECODE === "1";
+  const tile = q.shape[2]! > 1;
   if (tile && fusedSdpaSupported(q, mask, groupSize, bits)) {
     // The N-tiled loop CANNOT be traced: its tile count and slice extents
     // bake at trace-time N, so under shapeless replay a growing (concat-
-    // phase) quantized cache's newest rows are silently never attended —
-    // reproduced on e4b (whole-graph form, FUSED_DECODE=1: diverges within
-    // 3 tokens, then repetition-loops). generate() refuses to compile when
-    // MLX_BUN_FUSED_DECODE=1; this throw is the backstop for any other
-    // compiled caller — it lands in the trace's transactional fallback
-    // instead of producing silently wrong attention.
+    // phase) quantized cache's newest rows are silently never attended.
+    // This throw is the backstop for any compiled caller — it lands in the
+    // trace's transactional fallback instead of producing silently wrong
+    // attention.
     if (isCompiledTrace())
       throw new Error(
         "quantizedSdpaTiled inside a compiled trace: the tile loop bakes trace-time N " +