mlx-bun 0.0.13 → 0.3.0

package/src/model/flash-attention.ts

@@ -545,7 +545,7 @@ export function flashBackward(
     // (dK_j[d] → [kv_row][d]); no post-transpose. (An earlier "fix" transposed
     // it on the false belief the buffer was swapped — that transpose is only the
     // identity at Tkv==D and corrupted dK for Tkv≠D; removed after
-    // scripts/experiments/flash-dkv-debug.ts showed the raw output is correct.)
+    // flash-dkv-debug.ts (deleted 2026-08-23; git history) showed the raw output is correct.)
     return [dq[0]!, dkv[0]!, dkv[1]!];
   } finally {
     Dvec.dispose();
@@ -621,8 +621,8 @@ export function flashAttention(
 //  - This flash kernel is the L2 path: a port of mlx-optiq's
 //    flash_attention_metal (O(L) backward memory). Its parity oracle is OPTIQ,
 //    not finite differences. Confirmed: optiq's flash dK == ops.sdpa to f16
-//    (scripts/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
-//    (tests/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
+//    (scripts/oracle/flash-optiq-check.py, rel 0.0%), and mlx-bun's flash == ops.sdpa
+//    (tests/unit/flash-attention.test.ts) ⟹ mlx-bun flash == optiq flash (L2 parity).
 //
 // Two real port bugs were fixed to reach that parity (the original diverged
 // ~100% from optiq on dK): a spurious dK transpose in flashBackward (corrupted

package/src/model/gemma4-base.ts

@@ -1,5 +1,5 @@
 // gemma4-base — the config-INDEPENDENT machinery shared by every
-// architecture variant (docs/design/optimization_plan.md Phase B: pure code
+// architecture variant (docs/archive/investigations/optimization_plan.md Phase B: pure code
 // movement out of the gemma4.ts monolith; every definition here is
 // verbatim from it). Cache classes, quantized-KV SDPA, masks, the
 // quantized primitives + LoRA machinery, and small graph helpers live
@@ -13,8 +13,10 @@ import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } 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 { runtimeValue } from "../runtime-config";
 
 export type MaskMode = "" | "causal";
 export interface Mask {
@@ -53,6 +55,15 @@ export class LoraState {
   dropoutSeed: number | null = null;
 }
 
+/** Replaceable quantized payload. Inference sees this as read-only through
+ *  QuantizedLinear getters; sensitivity analysis exchanges it atomically. */
+export interface QuantizedLinearState {
+  w: MlxArray;
+  scales: MlxArray;
+  biases: MlxArray | null;
+  spec: ops.QuantSpec;
+}
+
 /** Inverted LoRA-input dropout, keyed by (seed, id) so it is deterministic —
  *  the same (seed, id, shape) reproduces the mask, which is what makes the
  *  segmented / gradient-checkpoint recompute correct. kept ⇒ x/(1-p),
@@ -79,18 +90,48 @@ export class QuantizedLinear {
   /** Stable per-target index for keying training-only LoRA dropout (set by
    *  attachForTraining). Gives each adapted linear an independent dropout mask. */
   dropoutId = 0;
+  private _w: MlxArray;
+  private _scales: MlxArray;
+  private _biases: MlxArray | null;
+  private _spec: ops.QuantSpec;
 
   constructor(
-    readonly w: MlxArray,
-    readonly scales: MlxArray,
-    readonly biases: MlxArray | null,
-    readonly spec: ops.QuantSpec,
+    w: MlxArray,
+    scales: MlxArray,
+    biases: MlxArray | null,
+    spec: ops.QuantSpec,
     /** ADDITIVE bias term (`.bias`, mlx nn.QuantizedLinear's optional bias —
      *  qwen2 qkv, starcoder2, …). Distinct from `biases`, the quantization
      *  zero-points. Applied after the matmul, before any LoRA residual,
      *  exactly like mlx's QuantizedLinear.__call__. */
     readonly bias: MlxArray | null = null,
-  ) {}
+  ) {
+    this._w = w;
+    this._scales = scales;
+    this._biases = biases;
+    this._spec = spec;
+  }
+
+  get w(): MlxArray { return this._w; }
+  get scales(): MlxArray { return this._scales; }
+  get biases(): MlxArray | null { return this._biases; }
+  get spec(): ops.QuantSpec { return this._spec; }
+
+  /** Atomically install a quantized payload and return the previous one.
+   *  This is the sole mutation capability used by sensitivity sweeps. */
+  exchangeQuantizedState(next: QuantizedLinearState): QuantizedLinearState {
+    const previous = {
+      w: this._w,
+      scales: this._scales,
+      biases: this._biases,
+      spec: this._spec,
+    };
+    this._w = next.w;
+    this._scales = next.scales;
+    this._biases = next.biases;
+    this._spec = next.spec;
+    return previous;
+  }
 
   static load(weights: Weights, path: string, config: ModelConfig): QuantizedLinear {
     if (!weights.has(`${path}.scales`))
@@ -109,14 +150,23 @@ export class QuantizedLinear {
   /** (in_features, out_features) from the quantized tensors (reference
    *  _infer_linear_shape: scales are [out, in/group_size]). */
   get inFeatures(): number {
-    return this.scales.shape[1]! * this.spec.groupSize;
+    return this._scales.shape[1]! * this._spec.groupSize;
   }
   get outFeatures(): number {
-    return this.scales.shape[0]!;
-  }
-
-  forward(x: MlxArray): MlxArray {
-    let out = ops.quantizedMatmul(x, this.w, this.scales, this.biases, this.spec, true);
+    return this._scales.shape[0]!;
+  }
+
+  forward(x: MlxArray, independentRows = false): MlxArray {
+    let out = independentRows ? quantizedMatmulRows(
+      x, this._w, this._scales, this._biases, this._spec,
+    ) : ops.quantizedMatmul(
+      x,
+      this._w,
+      this._scales,
+      this._biases,
+      this._spec,
+      true,
+    );
     // Additive bias (mlx QuantizedLinear: `x = x + self["bias"]`), before
     // the LoRA residual (LoRALinear calls the base linear bias-inclusive).
     if (this.bias) out = disposing(out, ops.add(out, this.bias));
@@ -216,6 +266,20 @@ export type SharedKv =
 
 export interface Cache {
   offset: number;
+  /** Stable storage identity for compatibility guards and persistence.
+   *  REQUIRED: every capability predicate (isPlainKvCache, isRotating*,
+   *  kv-store codecs) is a string compare on this value, so a missing
+   *  override used to fail OPEN into wrong routing (PRs #42/#43:
+   *  BatchedRotatingCache shipped without one and batched joins dropped
+   *  running rows). Wrappers report their own kind plus the inner kind. */
+  signature(): string;
+  /** Physical cache storage consumed by one additional token for one row.
+   *  Recurrent caches return 0 because their state does not grow with the
+   *  sequence. Optional for stateless/training adapters that are never
+   *  admitted or persisted. */
+  bytesPerToken?(): number;
+  /** state() returned temporary views that the caller must release. */
+  readonly stateNeedsDispose?: boolean;
   /** Compiled-decode trace adapters expose the RoPE offset as an int32
    *  array input here; real caches leave it unset (static int path). */
   readonly ropeOffsetArr?: MlxArray;
@@ -231,9 +295,56 @@ export interface Cache {
    *  speculative-decode rollback case (trim the last n ≤ γ tips right after a
    *  >1 concat write); rotating caches then physically shrink the buffer tail. */
   trim(n: number, bypass?: boolean): void;
+  /** Speculative verify-round support for NON-trimmable recurrent caches
+   *  (SSMCache — gated-DeltaNet conv + recurrent state). A cache exposing
+   *  these is spec-eligible without isTrimmable(): the serve loop arms a
+   *  round before the verify forward (the layer snapshots pre-round state —
+   *  free, MLX arrays are immutable — and records its position-local kernel
+   *  inputs), then resolves it: commit on full accept, or rollback(keep)
+   *  which restores the snapshot and bit-exactly REPLAYS the first `keep`
+   *  window tokens through the recurrence (identical arithmetic prefix ⇒
+   *  identical state). Trimmable caches leave these unset and keep trim(). */
+  specRoundBegin?(): void;
+  specRoundCommit?(): void;
+  specRoundRollback?(keep: number): void;
   dispose(): void;
 }
 
+export function cacheSignature(cache: Cache | undefined): string {
+  return cache ? cache.signature() : "unknown";
+}
+
+export interface RowBatchCache extends Cache {
+  readonly batchSize: number;
+  filterRows(keep: readonly number[]): void;
+  extractRow(row: number): Cache | null;
+}
+
+export function isRowBatchCache(cache: Cache): cache is RowBatchCache {
+  const candidate = cache as Partial<RowBatchCache>;
+  return typeof candidate.batchSize === "number" &&
+    typeof candidate.filterRows === "function" &&
+    typeof candidate.extractRow === "function";
+}
+
+export function isPlainKvCache(cache: Cache | undefined): cache is KVCache {
+  return cacheSignature(cache) === "kv:plain";
+}
+
+export function isQuantizedKvCache(cache: Cache | undefined): cache is QuantizedKVCache {
+  return cacheSignature(cache).startsWith("kv:quant:");
+}
+
+export function isRotatingPlainCache(cache: Cache | undefined): cache is RotatingKVCache {
+  return cacheSignature(cache) === "kv:rotating-plain";
+}
+
+export function isRotatingQuantizedCache(
+  cache: Cache | undefined,
+): cache is RotatingQuantizedKVCache {
+  return cacheSignature(cache).startsWith("kv:rotating-quant:");
+}
+
 /**
  * Capability contract for cache layouts that own their dynamic-row batching.
  * The scheduler uses this for non-K/V state families (for example GLM's
@@ -296,6 +407,13 @@ export class KVCache implements Cache {
   values: MlxArray | null = null;
   offset = 0;
 
+  signature(): string { return "kv:plain"; }
+
+  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]!;
@@ -469,6 +587,18 @@ export class QuantizedKVCache implements Cache {
 
   constructor(readonly groupSize: number, readonly bits: number) {}
 
+  signature(): string { return `kv:quant:${this.bits}:${this.groupSize}`; }
+
+  bytesPerToken(): number {
+    const steps = this.keys?.packed.shape[2] ?? 0;
+    if (steps === 0 || !this.keys || !this.values) return 0;
+    const tensors = [
+      this.keys.packed, this.keys.scales, this.keys.biases,
+      this.values.packed, this.values.scales, this.values.biases,
+    ];
+    return tensors.reduce((bytes, array) => bytes + array.nbytes, 0) / steps;
+  }
+
   updateAndFetch(): [MlxArray, MlxArray] {
     throw new Error("QuantizedKVCache: use updateAndFetchQuantized");
   }
@@ -681,6 +811,13 @@ export class RotatingKVCache implements Cache {
     this.maxSize = maxSize;
   }
 
+  signature(): string { return "kv:rotating-plain"; }
+
+  bytesPerToken(): number {
+    const steps = this.keys?.shape[2] ?? 0;
+    return steps > 0 ? (this.keys!.nbytes + this.values!.nbytes) / steps : 0;
+  }
+
   /** v with ring contents rearranged into temporal order (keep=0). */
   #temporalOrder(v: MlxArray): MlxArray {
     const S = v.shape[2]!;
@@ -1011,6 +1148,18 @@ export class RotatingQuantizedKVCache implements Cache {
     this.maxSize = maxSize;
   }
 
+  signature(): string { return `kv:rotating-quant:${this.bits}:${this.groupSize}`; }
+
+  bytesPerToken(): number {
+    const steps = this.keys?.packed.shape[2] ?? 0;
+    if (steps === 0 || !this.keys || !this.values) return 0;
+    const tensors = [
+      this.keys.packed, this.keys.scales, this.keys.biases,
+      this.values.packed, this.values.scales, this.values.biases,
+    ];
+    return tensors.reduce((bytes, array) => bytes + array.nbytes, 0) / steps;
+  }
+
   updateAndFetch(): [MlxArray, MlxArray] {
     throw new Error("RotatingQuantizedKVCache: use updateAndFetchQuantized");
   }
@@ -1358,7 +1507,7 @@ 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,
+ *  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 {
@@ -1377,7 +1526,7 @@ export const disposeTurboQuant = (t: TurboQuantTensor): void => {
   t.vScales.dispose();
 };
 
-/** TurboQuant KV cache — v1 (docs/design/turboquant-kv.md): dequantize-
+/** 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
  *  `#matches()` guard and every batching `instanceof` allow-list, so it
@@ -1396,6 +1545,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 {
+  readonly stateNeedsDispose = true;
   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
@@ -1406,9 +1556,22 @@ export class TurboQuantKVCache implements Cache {
   offset = 0;
   /** 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";
 
   constructor(readonly kBits: number, readonly vBits: number) {}
 
+  signature(): string { return `kv:turboquant:${this.kBits}:${this.vBits}`; }
+
+  bytesPerToken(): number {
+    const steps = this.#seqLen();
+    if (steps === 0 || !this.#kv) return 0;
+    return [
+      this.#kv.kIdx, this.#kv.kScales, this.#kv.kZeros,
+      this.#kv.vPacked, this.#kv.vScales,
+    ].reduce((bytes, array) => bytes + array.nbytes, 0) / steps;
+  }
+
   get headDim(): number | null {
     return this.#headDim;
   }
@@ -1473,6 +1636,13 @@ export class TurboQuantKVCache implements Cache {
       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();
@@ -1927,9 +2097,9 @@ export function quantizedSdpaTiled(
 function fusedSdpaSupported(q: MlxArray, mask: Mask, groupSize: number, bits: number): boolean {
   // Escape hatch mirroring optiq serve's --no-fused-kv: forces the
   // stock unfused path everywhere. Also the A/B lever for
-  // scripts/bench-fused-prefill.ts. Read per call (cheap next to the
+  // scripts/bench-levers.ts fused-prefill. Read per call (cheap next to the
   // FFI work) so tests and paired A/B harnesses can flip it in-process.
-  if (process.env.MLX_BUN_NO_FUSED_SDPA === "1") return false;
+  if (runtimeValue("MLX_BUN_NO_FUSED_SDPA") === "1") return false;
   if (bits !== 4 && bits !== 8) return false;
   if (groupSize !== 32 && groupSize !== 64 && groupSize !== 128) return false;
   if (q.dtype !== Dtype.bfloat16 && q.dtype !== Dtype.float16) return false;
@@ -1971,7 +2141,7 @@ export function quantizedSdpa(
  *  as a compile-time constant and call this for the rest. The combined
  *  predicate is exactly fusedSdpaSupported. */
 export function fusedSdpaRuntimeOk(q: MlxArray, mask: Mask): boolean {
-  if (process.env.MLX_BUN_NO_FUSED_SDPA === "1") return false;
+  if (runtimeValue("MLX_BUN_NO_FUSED_SDPA") === "1") return false;
   if (q.dtype !== Dtype.bfloat16 && q.dtype !== Dtype.float16) return false;
   if (mask.mode === "causal" || mask.mode === "") return true;
   if (mask.mode === "array")

package/src/model/gemma4.ts

@@ -23,7 +23,7 @@ import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
-import { isExpertTracing, recordRouting } from "../expert-trace";
+import { isExpertTracing, recordRouting } from "../lab/expert-trace/expert-trace";
 
 // The shared, config-independent machinery lives in gemma4-base.ts
 // (Phase B extraction); this file keeps the architecture-specific
@@ -56,7 +56,7 @@ 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 "../flags";
+import { flagOn } 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`
@@ -242,7 +242,7 @@ class Attention {
   /** A number offset takes the static fast::rope; an array offset (set by
    *  compiled-decode trace adapters) takes the dynamic variant — same
    *  kernel, offset read from the array (bit-exactness asserted in
-   *  tests/compile.test.ts). */
+   *  tests/unit/compile.test.ts). */
   rope(x: MlxArray, offset: number | MlxArray): MlxArray {
     return typeof offset === "number"
       ? ops.rope(x, this.headDim, this.ropeBase, offset, this.ropeFreqs)
@@ -370,7 +370,7 @@ class Attention {
       // verified route: its dQ/dK/dV vjps are exact (finite-difference checked;
       // it's what mlx-lm's tuner differentiates). The two historical bugs were
       // in THIS flash kernel (spurious dK transpose, divergent dQ barrier),
-      // both fixed and gated by tests/flash-attention.test.ts vs the ops.sdpa
+      // both fixed and gated by tests/unit/flash-attention.test.ts vs the ops.sdpa
       // oracle — see flash-attention.ts's routing comment. Flash also remains
       // Gemma-guarded in the trainer (e4b >=2K SIGTRAP not yet re-validated).
       // Full layers → window 0 (pure causal); sliding layers → their window.
@@ -885,7 +885,7 @@ export class Gemma4Model {
    *  which attend bidirectionally among themselves (use_bidirectional_attention:
    *  "vision" — text stays causal; audio prompts pass NO bidir mask, and a
    *  mixed image+audio prompt drops the image overlay too, §3.3 Q1 of
-   *  docs/design/audio-input-plan.md). `ids` ([1, L], the spliced token ids)
+   *  docs/design/generic-model-support.md). `ids` ([1, L], the spliced token ids)
    *  is required for per-layer-input models (e2b/e4b): multimodal soft-token
    *  positions get token 0's per-layer embedding, and their tower content
    *  enters only through the projection term (the merged hidden). Matches
@@ -1015,7 +1015,7 @@ export class Gemma4Model {
     return h;
   }
 
-  // --- Segmented-backward support (docs/design/segmented-backward-training.md
+  // --- 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
   // segment, threading the KV-shared donor K/V across segment boundaries.

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

@@ -4,14 +4,15 @@
 // fingerprint: 9f812d2eb461fcbe
 //
 // Branch-resolved, unrolled forward pass for this architecture
-// (docs/design/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
-// generator's gate: tests/generated-parity.test.ts. The cache-signature
+// (docs/archive/investigations/optimization_plan.md Phase C). Bit-exactness vs the monolith is the
+// generator's gate: tests/parity/generated-parity.test.ts. The cache-signature
 // guard falls back to the monolith for anything this file wasn't
 // generated for (bf16 compat runs, vision bidir prefill).
 
 import { MlxArray } from "../../mlx/array";
 import * as ops from "../../mlx/ops";
 import {
+  cacheSignature,
   disposing,
   fusedSdpaRuntimeOk,
   KVCache,
@@ -202,155 +203,58 @@ function donorFullQuantG64B4KeqVScal(layer: DecoderLayer, x: MlxArray, mask: Mas
 }
 
 export class GeneratedGemma4 extends Gemma4Model {
-  /** Per-layer cache classes must match what this file was generated
-   *  for; anything else runs the monolith path. */
+  /** Per-layer cache signatures must match what this file was generated for. */
   #matches(cache: Cache[]): boolean {
     return (
       cache.length === 48 &&
-      cache[0] instanceof RotatingQuantizedKVCache &&
-      cache[0].bits === 4 &&
-      cache[0].groupSize === 64 &&
-      cache[1] instanceof RotatingQuantizedKVCache &&
-      cache[1].bits === 4 &&
-      cache[1].groupSize === 64 &&
-      cache[2] instanceof RotatingQuantizedKVCache &&
-      cache[2].bits === 4 &&
-      cache[2].groupSize === 64 &&
-      cache[3] instanceof RotatingQuantizedKVCache &&
-      cache[3].bits === 8 &&
-      cache[3].groupSize === 64 &&
-      cache[4] instanceof RotatingQuantizedKVCache &&
-      cache[4].bits === 4 &&
-      cache[4].groupSize === 64 &&
-      cache[5] instanceof QuantizedKVCache &&
-      cache[5].bits === 4 &&
-      cache[5].groupSize === 64 &&
-      cache[6] instanceof RotatingQuantizedKVCache &&
-      cache[6].bits === 8 &&
-      cache[6].groupSize === 64 &&
-      cache[7] instanceof RotatingQuantizedKVCache &&
-      cache[7].bits === 8 &&
-      cache[7].groupSize === 64 &&
-      cache[8] instanceof RotatingQuantizedKVCache &&
-      cache[8].bits === 8 &&
-      cache[8].groupSize === 64 &&
-      cache[9] instanceof RotatingQuantizedKVCache &&
-      cache[9].bits === 4 &&
-      cache[9].groupSize === 64 &&
-      cache[10] instanceof RotatingQuantizedKVCache &&
-      cache[10].bits === 8 &&
-      cache[10].groupSize === 64 &&
-      cache[11] instanceof QuantizedKVCache &&
-      cache[11].bits === 4 &&
-      cache[11].groupSize === 64 &&
-      cache[12] instanceof RotatingQuantizedKVCache &&
-      cache[12].bits === 8 &&
-      cache[12].groupSize === 64 &&
-      cache[13] instanceof RotatingQuantizedKVCache &&
-      cache[13].bits === 4 &&
-      cache[13].groupSize === 64 &&
-      cache[14] instanceof RotatingQuantizedKVCache &&
-      cache[14].bits === 4 &&
-      cache[14].groupSize === 64 &&
-      cache[15] instanceof RotatingQuantizedKVCache &&
-      cache[15].bits === 4 &&
-      cache[15].groupSize === 64 &&
-      cache[16] instanceof RotatingQuantizedKVCache &&
-      cache[16].bits === 4 &&
-      cache[16].groupSize === 64 &&
-      cache[17] instanceof QuantizedKVCache &&
-      cache[17].bits === 4 &&
-      cache[17].groupSize === 64 &&
-      cache[18] instanceof RotatingQuantizedKVCache &&
-      cache[18].bits === 4 &&
-      cache[18].groupSize === 64 &&
-      cache[19] instanceof RotatingQuantizedKVCache &&
-      cache[19].bits === 4 &&
-      cache[19].groupSize === 64 &&
-      cache[20] instanceof RotatingQuantizedKVCache &&
-      cache[20].bits === 4 &&
-      cache[20].groupSize === 64 &&
-      cache[21] instanceof RotatingQuantizedKVCache &&
-      cache[21].bits === 4 &&
-      cache[21].groupSize === 64 &&
-      cache[22] instanceof RotatingQuantizedKVCache &&
-      cache[22].bits === 4 &&
-      cache[22].groupSize === 64 &&
-      cache[23] instanceof QuantizedKVCache &&
-      cache[23].bits === 4 &&
-      cache[23].groupSize === 64 &&
-      cache[24] instanceof RotatingQuantizedKVCache &&
-      cache[24].bits === 4 &&
-      cache[24].groupSize === 64 &&
-      cache[25] instanceof RotatingQuantizedKVCache &&
-      cache[25].bits === 4 &&
-      cache[25].groupSize === 64 &&
-      cache[26] instanceof RotatingQuantizedKVCache &&
-      cache[26].bits === 4 &&
-      cache[26].groupSize === 64 &&
-      cache[27] instanceof RotatingQuantizedKVCache &&
-      cache[27].bits === 4 &&
-      cache[27].groupSize === 64 &&
-      cache[28] instanceof RotatingQuantizedKVCache &&
-      cache[28].bits === 4 &&
-      cache[28].groupSize === 64 &&
-      cache[29] instanceof QuantizedKVCache &&
-      cache[29].bits === 4 &&
-      cache[29].groupSize === 64 &&
-      cache[30] instanceof RotatingQuantizedKVCache &&
-      cache[30].bits === 4 &&
-      cache[30].groupSize === 64 &&
-      cache[31] instanceof RotatingQuantizedKVCache &&
-      cache[31].bits === 4 &&
-      cache[31].groupSize === 64 &&
-      cache[32] instanceof RotatingQuantizedKVCache &&
-      cache[32].bits === 4 &&
-      cache[32].groupSize === 64 &&
-      cache[33] instanceof RotatingQuantizedKVCache &&
-      cache[33].bits === 4 &&
-      cache[33].groupSize === 64 &&
-      cache[34] instanceof RotatingQuantizedKVCache &&
-      cache[34].bits === 4 &&
-      cache[34].groupSize === 64 &&
-      cache[35] instanceof QuantizedKVCache &&
-      cache[35].bits === 4 &&
-      cache[35].groupSize === 64 &&
-      cache[36] instanceof RotatingQuantizedKVCache &&
-      cache[36].bits === 4 &&
-      cache[36].groupSize === 64 &&
-      cache[37] instanceof RotatingQuantizedKVCache &&
-      cache[37].bits === 4 &&
-      cache[37].groupSize === 64 &&
-      cache[38] instanceof RotatingQuantizedKVCache &&
-      cache[38].bits === 4 &&
-      cache[38].groupSize === 64 &&
-      cache[39] instanceof RotatingQuantizedKVCache &&
-      cache[39].bits === 4 &&
-      cache[39].groupSize === 64 &&
-      cache[40] instanceof RotatingQuantizedKVCache &&
-      cache[40].bits === 4 &&
-      cache[40].groupSize === 64 &&
-      cache[41] instanceof QuantizedKVCache &&
-      cache[41].bits === 4 &&
-      cache[41].groupSize === 64 &&
-      cache[42] instanceof RotatingQuantizedKVCache &&
-      cache[42].bits === 4 &&
-      cache[42].groupSize === 64 &&
-      cache[43] instanceof RotatingQuantizedKVCache &&
-      cache[43].bits === 4 &&
-      cache[43].groupSize === 64 &&
-      cache[44] instanceof RotatingQuantizedKVCache &&
-      cache[44].bits === 4 &&
-      cache[44].groupSize === 64 &&
-      cache[45] instanceof RotatingQuantizedKVCache &&
-      cache[45].bits === 4 &&
-      cache[45].groupSize === 64 &&
-      cache[46] instanceof RotatingQuantizedKVCache &&
-      cache[46].bits === 4 &&
-      cache[46].groupSize === 64 &&
-      cache[47] instanceof QuantizedKVCache &&
-      cache[47].bits === 4 &&
-      cache[47].groupSize === 64
+      cacheSignature(cache[0]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[1]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[2]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[3]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[4]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[5]) === "kv:quant:4:64" &&
+      cacheSignature(cache[6]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[7]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[8]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[9]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[10]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[11]) === "kv:quant:4:64" &&
+      cacheSignature(cache[12]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[13]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[14]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[15]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[16]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[17]) === "kv:quant:4:64" &&
+      cacheSignature(cache[18]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[19]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[20]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[21]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[22]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[23]) === "kv:quant:4:64" &&
+      cacheSignature(cache[24]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[25]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[26]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[27]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[28]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[29]) === "kv:quant:4:64" &&
+      cacheSignature(cache[30]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[31]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[32]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[33]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[34]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[35]) === "kv:quant:4:64" &&
+      cacheSignature(cache[36]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[37]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[38]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[39]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[40]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[41]) === "kv:quant:4:64" &&
+      cacheSignature(cache[42]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[43]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[44]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[45]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[46]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[47]) === "kv:quant:4:64"
     );
   }