mlx-bun 0.0.12 → 0.2.0

package/src/model/gemma4-base.ts

@@ -15,6 +15,7 @@ import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
 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 +54,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 +89,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 +149,21 @@ 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]!;
+    return this._scales.shape[0]!;
   }
 
   forward(x: MlxArray): MlxArray {
-    let out = ops.quantizedMatmul(x, this.w, this.scales, this.biases, this.spec, true);
+    let out = 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 +263,15 @@ export type SharedKv =
 
 export interface Cache {
   offset: number;
+  /** Stable storage identity for compatibility guards and persistence. */
+  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 +287,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?.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 +399,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 +579,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 +803,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 +1140,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");
   }
@@ -1396,6 +1537,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
@@ -1409,6 +1551,17 @@ export class TurboQuantKVCache implements Cache {
 
   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;
   }
@@ -1929,7 +2082,7 @@ function fusedSdpaSupported(q: MlxArray, mask: Mask, groupSize: number, bits: nu
   // stock unfused path everywhere. Also the A/B lever for
   // scripts/bench-fused-prefill.ts. 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 +2124,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/generated/gemma4-12b.ts

@@ -12,6 +12,7 @@
 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"
     );
   }
 

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

@@ -12,6 +12,7 @@
 import { MlxArray } from "../../mlx/array";
 import * as ops from "../../mlx/ops";
 import {
+  cacheSignature,
   disposing,
   fusedSdpaRuntimeOk,
   KVCache,
@@ -293,101 +294,40 @@ function donorFullQuantG64B4KeqVMoeScal(layer: DecoderLayer, x: MlxArray, mask:
 }
 
 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 === 30 &&
-      cache[0] instanceof RotatingQuantizedKVCache &&
-      cache[0].bits === 8 &&
-      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 === 8 &&
-      cache[4].groupSize === 64 &&
-      cache[5] instanceof QuantizedKVCache &&
-      cache[5].bits === 8 &&
-      cache[5].groupSize === 64 &&
-      cache[6] instanceof RotatingQuantizedKVCache &&
-      cache[6].bits === 4 &&
-      cache[6].groupSize === 64 &&
-      cache[7] instanceof RotatingQuantizedKVCache &&
-      cache[7].bits === 4 &&
-      cache[7].groupSize === 64 &&
-      cache[8] instanceof RotatingQuantizedKVCache &&
-      cache[8].bits === 8 &&
-      cache[8].groupSize === 64 &&
-      cache[9] instanceof RotatingQuantizedKVCache &&
-      cache[9].bits === 8 &&
-      cache[9].groupSize === 64 &&
-      cache[10] instanceof RotatingQuantizedKVCache &&
-      cache[10].bits === 4 &&
-      cache[10].groupSize === 64 &&
-      cache[11] instanceof QuantizedKVCache &&
-      cache[11].bits === 8 &&
-      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
+      cacheSignature(cache[0]) === "kv:rotating-quant:8: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:8:64" &&
+      cacheSignature(cache[5]) === "kv:quant:8:64" &&
+      cacheSignature(cache[6]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[7]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[8]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[9]) === "kv:rotating-quant:8:64" &&
+      cacheSignature(cache[10]) === "kv:rotating-quant:4:64" &&
+      cacheSignature(cache[11]) === "kv:quant:8: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"
     );
   }