@mlx-bun/inference 0.0.0-stage → 1.0.0

package/src/state/grouped-heads-kv4-cache.ts

@@ -0,0 +1,26 @@
+import { quantizedSdpaGroupedHeads } from "../layers/quantized-attention";
+import { unfusedKernel, type AffineKernels } from "./affine-attention";
+import { AttentionMasks } from "./attention-read";
+import { BatchedQuantizedKVCache } from "./batched-quantized-kv";
+
+const grouped: AffineKernels["window"] = (q, keys, values, scale, mask) =>
+  quantizedSdpaGroupedHeads(q, keys, values, scale, mask, 64, 4);
+const kernels: AffineKernels = { decode: unfusedKernel(64, 4), window: grouped, maskedWindow: grouped };
+
+/** The 4-bit group-64 batched full-attention cache of a depth-2 speculative
+ * verify composition on the M4 Pro GPU (applegpu_g16s): `BatchedQuantizedKVCache`
+ * whose window read is the grouped-heads key matmul
+ * (`quantizedSdpaGroupedHeads`, three GQA heads per batch, where MLX's native
+ * key matvec keeps exact arithmetic at 8192 or more keys) and whose decode read
+ * is the stock unfused port. Masks are the batched layout's.
+ *
+ * The window kernel's contract is the verify window: 3 queries per row, 24
+ * query heads, 4 KV heads, head dim 256, B <= 2, bf16 or f32 queries. Today
+ * `quantizedSdpaUnfused` applied it on that GPU at those shapes and 8192 or
+ * more keys, whenever the dispatch ran unfused; this cache applies it to every
+ * window, so it serves only a composition whose windows are all verify windows.
+ * Constructed explicitly; nothing selects it by device. */
+export class GroupedHeadsKv4Cache extends BatchedQuantizedKVCache {
+  constructor(masks = new AttentionMasks()) { super(64, 4, kernels, masks); }
+  override makeEmptyBatch(): GroupedHeadsKv4Cache { return new GroupedHeadsKv4Cache(this.masks); }
+}

package/src/state/index.ts

@@ -0,0 +1,38 @@
+export * from "./kv";
+export * from "./quantized-kv";
+export { composedAffineKernels, tiledAffineKernels, tiledCausalAffineKernels, unfusedAffineKernels, type AffineKernels } from "./affine-attention";
+export * from "./rotating-kv";
+export * from "./rotating-quantized-kv";
+export * from "./turboquant-kv";
+export { deferredTurboQuantKernels, denseTurboQuantKernels, uncomposedTurboQuantKernels, type TurboQuantKernels } from "./turboquant-attention";
+export * from "./turboquant-codec";
+export * from "./ssm";
+export * from "./glm52-cache";
+export * from "./views";
+export * from "./checkpoint";
+export { createKvMaintenance, unchangedKv, type KvMaintenance } from "./layout";
+export * from "./kv-scheme";
+export * from "./persistence";
+export * from "./batched-kv";
+export * from "./batched-quantized-kv";
+export * from "./batched-turboquant-kv";
+export * from "./batched-rotating";
+export * from "./batched-rotating-quant";
+export * from "./bf16-first-kv";
+export * from "./delayed-quantized-kv";
+export * from "./delayed-turboquant-kv";
+export * from "./delayed-rotating-quantized-kv";
+export * from "./batched-ssm";
+export * from "./speculative-rotating-kv";
+export * from "./rollback";
+export * from "./encoder-cache";
+export * from "./prefix-cache";
+export * from "./retention-policy";
+export * from "./rows";
+export * from "./attachment";
+export * from "./ssd-cache";
+export * from "./tiered-prefix-cache";
+export * from "./durability";
+export * from "./continuation-persistence";
+export * from "./kv-budget";
+export * from "./cache-service";

package/src/state/kv-attention-view.ts

@@ -0,0 +1,77 @@
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import * as ops from "@mlx-bun/mlx/ops";
+import type { Cache, KvAttentionView, KvDonorAttention } from "../contracts/mlx/cache";
+import { disposeTriple } from "./quantized-tensor";
+import { kvStorage } from "./kv-storage";
+import { quantizedAppendAttention } from "../layers/quantized-append-attention";
+import { viewAffineKernel, type AffineKernels } from "./affine-attention";
+
+export function captureKvDonorAttention(cache: Cache): KvDonorAttention {
+  if (cache.captureDonorAttention) return cache.captureDonorAttention();
+  const view = cache.captureDonorRows!();
+  return {
+    width: view.keys.shape[2]!, dtype: view.keys.dtype,
+    offsets: view.offsets, starts: view.starts, ends: view.ends,
+    attend: (q, scale, mask) => ops.sdpa(q, view.keys, view.values, scale, mask.mode, mask.arr),
+    dispose() { view.keys.dispose(); view.values.dispose(); },
+  };
+}
+
+/** Q-only attention over fetched affine keys and values with the storage's
+ * kernel set (`viewAffineKernel`). Takes ownership of `keys` and `values`. */
+export function quantizedDonorAttention(keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
+  validity: Pick<KvDonorAttention, "offsets" | "starts" | "ends">, kernels: AffineKernels): KvDonorAttention {
+  return {
+    width: keys.packed.shape[2]!, dtype: keys.scales.dtype, ...validity,
+    attend: (q, scale, mask) => viewAffineKernel(kernels, q.shape[2]!, mask)(q, keys, values, scale, mask),
+    dispose() { disposeTriple(keys); disposeTriple(values); },
+  };
+}
+
+/** Append `k` and `v` to `cache` and return its attention view. Plain storage
+ * reads with the stock SDPA; affine storage with `kernels`, the kernel set of
+ * the row layout reading it (`viewAffineKernel`). */
+export function captureKvAttention(row: Cache, k: MlxArray, v: MlxArray, kernels: AffineKernels): KvAttentionView {
+    const cache = kvStorage(row);
+    if (cache.attentionState) return cache.attentionState.appendAndFetch(k, v);
+    const quantized = cache.quantizedAttention;
+    if (quantized) {
+      const { groupSize, bits } = quantized;
+      const [keys, values] = quantized.updateAndFetchQuantized(k, v);
+      return {
+        attend: (q, scale, mask, independentPositions) => independentPositions && q.shape[2]! > 1
+          ? quantizedAppendAttention(q, keys, values, scale, groupSize, bits)
+          : viewAffineKernel(kernels, q.shape[2]!, mask)(q, keys, values, scale, mask),
+        dispose() { disposeTriple(keys); disposeTriple(values); },
+      };
+    }
+    const [keys, values] = cache.updateAndFetch(k, v);
+    return {
+      attend: (q, scale, mask) => ops.sdpa(q, keys, values, scale, mask.mode, mask.arr),
+      dispose() { keys.dispose(); values.dispose(); },
+    };
+  }
+
+
+/** Takes ownership of already aligned per-row donor views. Each storage codec
+ * keeps its own attention arithmetic while the caller supplies one query batch. */
+export function combineKvDonorAttention(rows: readonly KvDonorAttention[]): KvDonorAttention {
+  return { width: rows[0]!.width, dtype: rows[0]!.dtype,
+    offsets: rows.flatMap(row => [...row.offsets]), starts: rows.flatMap(row => [...row.starts]),
+    ends: rows.flatMap(row => [...row.ends]),
+    attend(q, scale, mask) {
+      const outputs: MlxArray[] = [];
+      try {
+        for (const [index, row] of rows.entries()) {
+          using query = q.slice([index, 0, 0, 0], [index + 1, q.shape[1]!, q.shape[2]!, q.shape[3]!]);
+          const from = mask.arr?.shape.map(() => 0), to = mask.arr ? [...mask.arr.shape] : undefined;
+          if (from && to && to.length === 4 && to[0]! > 1) { from[0] = index; to[0] = index + 1; }
+          using arr = mask.arr ? mask.arr.slice(from!, to!) : null;
+          outputs.push(row.attend(query, scale, { mode: mask.mode, arr }));
+        }
+        return outputs.length === 1 ? ops.contiguous(outputs[0]!) : ops.concatAxis(outputs, 0);
+      } finally { for (const output of outputs) output.dispose(); }
+    },
+    dispose() { for (const row of rows) row.dispose(); },
+  };
+}

package/src/state/kv-budget.ts

@@ -0,0 +1,14 @@
+import type { ModelConfig } from "../artifacts/config";
+import { kvBytesAt, KvScheme } from "./kv-scheme";
+
+/** Project one batch row's worst-case KV bytes using the same per-layer
+ * scheme that the server-wide admission ceiling uses. */
+export function batchRowKvBytes(
+  config: ModelConfig,
+  promptTokens: number,
+  maxTokens: number,
+  scheme?: KvScheme,
+): number {
+  const tokens = promptTokens + maxTokens;
+  return scheme ? scheme.bytesAt(config, tokens) : kvBytesAt(config, tokens);
+}

package/src/state/kv-scheme.ts

@@ -0,0 +1,315 @@
+import type { KvQuantSpec, ModelConfig, TurboQuantScheme } from "../artifacts/config";
+import type { AffineKernelChoice, KvStorageScheme } from "../contracts/portable/composition";
+
+const BF16_BYTES = 2;
+
+export type KvQuantOverride = "off" | "config" | number | undefined;
+
+export interface KvSchemeOptions {
+  kvBits?: number;
+  kvGroupSize?: number;
+  kvConfig?: KvQuantSpec[];
+  quantizedKvStart?: number;
+  turboQuant?: TurboQuantScheme;
+}
+
+interface ResolvedKvSchemeOptions extends Omit<KvSchemeOptions, "kvConfig" | "turboQuant"> {
+  readonly kvConfig?: readonly Readonly<KvQuantSpec>[];
+  readonly turboQuant?: Readonly<TurboQuantScheme>;
+}
+
+export type KvSchemeKind = "bf16" | "affine-uniform" | "affine-config" | "turbo";
+
+/** How one attention layer stores its cache: every position (`full`), a rolling
+ * window (`sliding`), or a fixed-size recurrent state with no KV (`recurrent`). */
+export type CacheLayerKind = "full" | "sliding" | "recurrent";
+
+/** The cache layout the config declares for a layer. A layer the config does not
+ * list is full attention. The one place a layer-type spelling is read. */
+export function cacheLayerKind(config: Pick<ModelConfig, "text">, layer: number): CacheLayerKind {
+  const type = config.text.layerTypes[layer];
+  return type === "linear_attention" ? "recurrent" : type === "sliding_attention" ? "sliding" : "full";
+}
+
+export interface KvGeometry {
+  fullBytesPerToken: number;
+  slidingBytesPerToken: number;
+  linearStateBytes: number;
+  window: number;
+}
+
+/** Bytes per stored affine-quantized KV element, including scale and bias. */
+export function kvQuantBytesPerElement(bits: number, groupSize: number): number {
+  return bits / 8 + 4 / groupSize;
+}
+
+export function kvGeometry(config: ModelConfig, options: KvSchemeOptions = {}): KvGeometry {
+  const text = config.text;
+  const byLayer = options.kvConfig?.length
+    ? new Map(options.kvConfig.map((entry) => [entry.layerIdx, entry]))
+    : null;
+  const uniformBytes = options.kvBits
+    ? kvQuantBytesPerElement(options.kvBits, options.kvGroupSize ?? 64)
+    : BF16_BYTES;
+  const elementBytes = (layer: number): number => {
+    const entry = byLayer?.get(layer);
+    return entry
+      ? kvQuantBytesPerElement(entry.bits, entry.groupSize)
+      : byLayer ? BF16_BYTES : uniformBytes;
+  };
+
+  let linearLayers = 0;
+  let fullBytesPerToken = 0;
+  let slidingBytesPerToken = 0;
+  for (let layer = 0; layer < text.numHiddenLayers; layer++) {
+    const kind = cacheLayerKind(config, layer);
+    if (kind === "recurrent") {
+      linearLayers++;
+    } else if (kind === "sliding") {
+      slidingBytesPerToken +=
+        2 * text.numKeyValueHeads * text.headDim * elementBytes(layer);
+    } else {
+      fullBytesPerToken +=
+        2 * text.numGlobalKeyValueHeads * text.globalHeadDim * elementBytes(layer);
+    }
+  }
+  const convDim = 2 * text.linearNumKeyHeads * text.linearKeyHeadDim +
+    text.linearNumValueHeads * text.linearValueHeadDim;
+  const linearStateBytes = linearLayers === 0 ? 0 : linearLayers * (
+    text.linearNumValueHeads * text.linearValueHeadDim * text.linearKeyHeadDim * 4 +
+    Math.max(0, text.linearConvKernelDim - 1) * convDim * BF16_BYTES
+  );
+  return {
+    fullBytesPerToken,
+    slidingBytesPerToken,
+    linearStateBytes,
+    window: text.slidingWindow,
+  };
+}
+
+export function kvBytesAt(
+  config: ModelConfig,
+  tokens: number,
+  options: KvSchemeOptions = {},
+): number {
+  const value = kvGeometry(config, options);
+  return value.fullBytesPerToken * tokens +
+    value.slidingBytesPerToken * Math.min(tokens, value.window) +
+    value.linearStateBytes;
+}
+
+/** Head dims MLX 0.32.2's fused multi-query SDPA kernel serves for L > 8
+ *  queries on M1–M4 GPUs (mlx/backend/metal/scaled_dot_product_attention.cpp
+ *  `use_fallback`: 192 and 256 always fall back off-NAX; `has_fused_kernel`
+ *  lists 64/72/80/96/128 for the full kernel). Every other head dim runs
+ *  mlx/fast.cpp's composed fallback, which materializes the bf16 scores
+ *  tensor [heads, L, S] plus a bool mask per full-attention layer. The
+ *  decode kernel (L ≤ 8) fuses 256, so this term is a prefill-only cost. */
+export const FUSED_SDPA_HEAD_DIMS: ReadonlySet<number> = new Set([64, 72, 80, 96, 128]);
+const SDPA_SCORES_BYTES = 2; // bf16 scores in q's dtype
+const SDPA_FUSED_MAX_QUERIES = 8;
+
+/** Peak bytes of the SDPA-fallback scores tensor live during one prefill
+ *  chunk of `chunk` queries against a `ctx`-token context: the last chunk
+ *  sees the longest keys, and the lazy graph keeps one layer's scores live
+ *  at a time. Zero when every attention layer's head dim fuses. Sliding
+ *  layers see at most `window + L` keys. Admission-only (fit and the
+ *  in-flight memory guard); nothing here changes what is computed. */
+export function sdpaFallbackBytes(config: ModelConfig, chunk: number, ctx: number): number {
+  const text = config.text;
+  const heads = text.numAttentionHeads || 0;
+  const queries = Math.min(chunk, ctx);
+  if (heads <= 0 || queries <= SDPA_FUSED_MAX_QUERIES) return 0;
+  let peak = 0;
+  const layers = text.numHiddenLayers || text.layerTypes.length;
+  for (let layer = 0; layer < layers; layer++) {
+    const kind = cacheLayerKind(config, layer);
+    if (kind === "recurrent") continue;
+    const sliding = kind === "sliding" && text.slidingWindow > 0;
+    const headDim = sliding ? text.headDim : text.globalHeadDim;
+    if (FUSED_SDPA_HEAD_DIMS.has(headDim)) continue;
+    const keys = sliding ? Math.min(ctx, text.slidingWindow + queries) : ctx;
+    peak = Math.max(peak, heads * queries * keys * SDPA_SCORES_BYTES);
+  }
+  return peak;
+}
+
+/** The conversion start a scheme's caches keep bf16 for: the explicit
+ * `quantizedKvStart`, else immediately for a per-layer config and TurboQuant
+ * and after 5000 tokens for uniform bits (mlx-lm's default). */
+function conversionStart(kind: KvSchemeKind, options: Readonly<ResolvedKvSchemeOptions>): number {
+  return options.quantizedKvStart ?? (kind === "affine-uniform" ? 5000 : 0);
+}
+
+/** The composition record's form of a scheme: what the graph's `makeCache`
+ * composes, with the affine read kernels and the conversion start (carried
+ * only when nonzero). */
+function storageOf(kind: KvSchemeKind, options: Readonly<ResolvedKvSchemeOptions>, kernels: AffineKernelChoice): KvStorageScheme {
+  const tokens = conversionStart(kind, options);
+  const start = tokens > 0 ? { quantizedKvStart: tokens } : {};
+  switch (kind) {
+    case "bf16": return Object.freeze({ kind: "bf16" });
+    case "affine-uniform":
+      return Object.freeze({ kind: "affine", bits: options.kvBits!, groupSize: options.kvGroupSize ?? 64, ...start, kernels });
+    case "affine-config":
+      return Object.freeze({ kind: "affine-layers", layers: Object.freeze(options.kvConfig!.map(entry =>
+        Object.freeze({ layer: entry.layerIdx, bits: entry.bits, groupSize: entry.groupSize }))), ...start, kernels });
+    case "turbo":
+      return Object.freeze({ kind: "turbo", kBits: options.turboQuant!.kBits, vBits: options.turboQuant!.vBits, ...start });
+  }
+}
+
+/** The kernels a scheme composes when `--fused-sdpa` is unset: OptiQ's
+ * tiled-causal rule for a per-layer config, mlx-lm's unfused port otherwise. */
+function defaultKernels(kind: KvSchemeKind): AffineKernelChoice {
+  return kind === "affine-config" ? "tiled-causal" : "unfused";
+}
+
+export class KvScheme {
+  readonly kind: KvSchemeKind;
+  readonly options: Readonly<ResolvedKvSchemeOptions>;
+  /** The composition record's form (`Composition.kv`): the storage every
+   * attention cache of the graph composes and the read kernels of its affine
+   * caches. */
+  readonly storage: KvStorageScheme;
+
+  constructor(kind: KvSchemeKind, options: KvSchemeOptions, kernels: AffineKernelChoice = defaultKernels(kind)) {
+    this.kind = kind;
+    const kvConfig = options.kvConfig?.map((entry) =>
+      Object.freeze({ ...entry }),
+    );
+    const turboQuant = options.turboQuant
+      ? Object.freeze({ ...options.turboQuant })
+      : undefined;
+    this.options = Object.freeze({
+      ...options,
+      ...(kvConfig ? { kvConfig: Object.freeze(kvConfig) } : {}),
+      ...(turboQuant ? { turboQuant } : {}),
+    });
+    this.storage = storageOf(kind, this.options, kernels);
+  }
+
+  get quantized(): boolean {
+    return this.kind !== "bf16";
+  }
+
+  get cacheKey(): string {
+    if (this.kind === "turbo") {
+      const turbo = this.options.turboQuant!;
+      const start = this.options.quantizedKvStart ?? 0;
+      return `turbo-k${turbo.kBits}v${turbo.vBits}${start > 0 ? `-start${start}` : ""}`;
+    }
+    if (this.kind === "affine-uniform") {
+      const start = this.options.quantizedKvStart ?? 5000, group = this.options.kvGroupSize ?? 64;
+      return `kv${this.options.kvBits}${group !== 64 ? `-g${group}` : ""}${start > 0 ? `-start${start}` : ""}`;
+    }
+    if (this.kind === "affine-config") {
+      const start = this.options.quantizedKvStart ?? 0;
+      return `config${start > 0 ? `-start${start}` : ""}`;
+    }
+    return "bf16";
+  }
+
+  get label(): string {
+    if (this.kind === "turbo") {
+      const turbo = this.options.turboQuant!;
+      return `turbo k${turbo.kBits}v${turbo.vBits}`;
+    }
+    if (this.kind === "affine-uniform") return `uniform-kv${this.options.kvBits}`;
+    if (this.kind === "affine-config") return "mixed (kv_config.json)";
+    return "bf16";
+  }
+
+  /** Configured attention storage, not live allocator occupancy. Recurrent
+   * layers retain their own state and never acquire an attention KV codec. */
+  describe(config: ModelConfig) {
+    const layers: Record<string, number> = {};
+    const attention = { global: 0, sliding_window: 0 };
+    let recurrent_layers = 0;
+    const perLayer = new Map(this.options.kvConfig?.map(entry => [entry.layerIdx, entry.bits]));
+    for (let layer = 0; layer < config.text.numHiddenLayers; layer++) {
+      const kind = cacheLayerKind(config, layer);
+      if (kind === "recurrent") { recurrent_layers++; continue; }
+      const sliding = kind === "sliding";
+      attention[sliding ? "sliding_window" : "global"]++;
+      const bits = this.kind === "affine-config" ? perLayer.get(layer)
+        : this.kind === "affine-uniform" ? this.options.kvBits : undefined;
+      const codec = this.kind === "turbo" && !sliding
+        ? `turbo-k${this.options.turboQuant!.kBits}v${this.options.turboQuant!.vBits}`
+        : bits ? `kv${bits}` : "bf16";
+      layers[codec] = (layers[codec] ?? 0) + 1;
+    }
+    return { mode: this.label, layers, attention, recurrent_layers };
+  }
+
+  bytesAt(config: ModelConfig, tokens: number): number {
+    // TurboQuant stays conservatively billed as bf16 until its packed layout
+    // exposes a stable projector.
+    return kvBytesAt(config, tokens, this.fitOptions);
+  }
+
+  /** Owned copies of the options for memory fit; TurboQuant fits as bf16. */
+  get fitOptions(): KvSchemeOptions {
+    if (this.kind === "turbo") return {};
+    const { kvConfig, turboQuant: _turbo, ...scalar } = this.options;
+    return { ...scalar, ...(kvConfig ? { kvConfig: kvConfig.map((entry) => ({ ...entry })) } : {}) };
+  }
+}
+
+/** Resolve the KV flags (`--kv-quant`, its TurboQuant spelling,
+ * `--quantized-kv-start`, `--fused-sdpa`) against the artifact's
+ * `kv_config.json`. `fusedSdpa` picks the affine read kernels (on:
+ * tiled-causal, off: unfused); unset, they follow the scheme's kind. */
+export function resolveKvScheme(input: {
+  override?: KvQuantOverride;
+  turboQuant?: TurboQuantScheme;
+  config?: readonly KvQuantSpec[] | null;
+  missingConfig?: "bf16" | "error";
+  quantizedKvStart?: number;
+  fusedSdpa?: boolean;
+}): KvScheme {
+  const start = input.quantizedKvStart;
+  if (start !== undefined && (!Number.isSafeInteger(start) || start < 0))
+    throw new Error("quantizedKvStart must be a nonnegative integer");
+  const kernels = (kind: KvSchemeKind): AffineKernelChoice =>
+    input.fusedSdpa === undefined ? defaultKernels(kind) : input.fusedSdpa ? "tiled-causal" : "unfused";
+  if (input.turboQuant) {
+    return new KvScheme("turbo", {
+      turboQuant: input.turboQuant,
+      quantizedKvStart: start ?? 0,
+    });
+  }
+  if (input.override === "off" || input.override === undefined)
+    return new KvScheme("bf16", {});
+  if (input.override === "config") {
+    if (!input.config?.length) {
+      if (input.missingConfig === "error")
+        throw new Error("model has no kv_config.json (--kv-quant config)");
+      return new KvScheme("bf16", {});
+    }
+    return new KvScheme("affine-config", { kvConfig: [...input.config],
+      ...(start === undefined ? {} : { quantizedKvStart: start }) }, kernels("affine-config"));
+  }
+  return new KvScheme("affine-uniform", {
+    kvBits: input.override,
+    quantizedKvStart: start ?? 0,
+  }, kernels("affine-uniform"));
+}
+
+/** The scheme a composition record names (`Composition.kv`): the inverse of
+ * `KvScheme.storage`, for what describes, bills and keys the composed caches. */
+export function kvSchemeOf(kv: KvStorageScheme): KvScheme {
+  switch (kv.kind) {
+    case "bf16": return new KvScheme("bf16", {});
+    case "affine":
+      return new KvScheme("affine-uniform", { kvBits: kv.bits, kvGroupSize: kv.groupSize,
+        quantizedKvStart: kv.quantizedKvStart ?? 0 }, kv.kernels);
+    case "affine-layers":
+      return new KvScheme("affine-config", { kvConfig: kv.layers.map(entry =>
+        ({ layerIdx: entry.layer, bits: entry.bits, groupSize: entry.groupSize })),
+        quantizedKvStart: kv.quantizedKvStart ?? 0 }, kv.kernels);
+    case "turbo":
+      return new KvScheme("turbo", { turboQuant: { kBits: kv.kBits, vBits: kv.vBits },
+        quantizedKvStart: kv.quantizedKvStart ?? 0 });
+  }
+}

package/src/state/kv-storage.ts

@@ -0,0 +1,57 @@
+// The storage-level reads and conversions state's own caches use on one
+// another: a row layout fetching its serial rows' keys and values, a delayed
+// lego converting the storage it holds, a donor capture reading packed or
+// rotated storage. They are internal to `state/`. Graphs, layers and every
+// other caller read a cache through the cache contract (`AttentionCache` and
+// its named reads, `GatedDeltaCache`), which carries none of these.
+
+import type { MlxArray } from "@mlx-bun/mlx/array";
+import type { Dtype } from "@mlx-bun/mlx/ffi";
+import type * as ops from "@mlx-bun/mlx/ops";
+import type { Cache, KvAttentionView, Mask } from "../contracts/mlx/cache";
+
+/** Affine-quantized storage: appends and hands back its packed keys and values. */
+export interface QuantizedAttentionState {
+  readonly groupSize: number;
+  readonly bits: number;
+  updateAndFetchQuantized(k: MlxArray, v: MlxArray): [ops.QuantizedTensor, ops.QuantizedTensor];
+}
+
+/** TurboQuant storage whose values 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];
+}
+
+/** Storage that appends once and hands every attention consumer the same view. */
+export interface KvAttentionState {
+  appendAndFetch(k: MlxArray, v: MlxArray): KvAttentionView;
+}
+
+/** A cache as `state/` sees its storage. */
+export interface KvStorage extends Cache {
+  /** Appends `k` and `v` and returns the stored keys and values (owned views). */
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray];
+  /** The mask an `N`-position append would read under, from this storage's state. */
+  makeMask(N: number, windowSize: number | null): Mask;
+  /** Affine conversion of this storage. Logical position can differ from the
+   * physical write head in a padded rotating layout. `kernels` builds the
+   * kernel set the converted storage attends with (`AffineKernelSet`) for its
+   * bits, group size and activation dtype. */
+  readonly affineConversion?: { readonly offset: number;
+    toQuantized(groupSize: number, bits: number, kernels: (bits: number, groupSize: number, dtype: Dtype) => unknown): Cache };
+  /** TurboQuant conversion of this storage. */
+  readonly turboConversion?: { readonly offset: number; toTurboQuantized(kBits: number, vBits: number): Cache };
+  readonly attentionState?: KvAttentionState;
+  readonly rotatedValueAttention?: RotatedValueAttentionState;
+  readonly quantizedAttention?: QuantizedAttentionState;
+}
+
+/** `cache` as storage, for a caller inside `state/` holding a KV storage
+ * (every key/value storage here implements these reads; the contract does not
+ * promise them, so nothing outside `state/` may call them). */
+export function kvStorage(cache: Cache): KvStorage {
+  return cache as KvStorage;
+}