mlx-bun 0.0.13 → 0.3.0

package/src/model/profile.ts

@@ -0,0 +1,443 @@
+import type { ModelConfig } from "../config";
+import { configFingerprint } from "./fingerprint";
+import {
+  isDiffusionGemmaConfig,
+  isGlm52Config,
+  isMiniCPM5Config,
+  isQwen35Config,
+  isQwen3Config,
+  isQwen3MoeConfig,
+} from "./support";
+import { GENERATED } from "./generated";
+import { GENERIC_MODEL_TYPES, genericArgsFor, remapModelType } from "./universal/archs";
+
+export const ENGINE_CAPABILITIES = Object.freeze([
+  "autoregressive",
+  "colibri-container",
+  "diffusion",
+  "diffusion-gemma-graph",
+  "gemma4-graph",
+  "generated-graph",
+  "glm5.2-graph",
+  "minicpm5-graph",
+  "mixed-precision-kv",
+  "native-mtp",
+  "qwen3-graph",
+  "qwen3-moe-graph",
+  "qwen3.5-graph",
+  "recurrent-state",
+  "safetensors",
+  "streamed-experts",
+  "universal-dense-graph",
+  "vision-sidecar",
+] as const);
+
+export type EngineCapability = typeof ENGINE_CAPABILITIES[number];
+export type FidelityTier = "l1" | "l2" | "l3";
+
+export type FidelityTarget =
+  | Readonly<{ tier: "l1"; oracle: "mlx-lm"; claim: "bit-exact" }>
+  | Readonly<{ tier: "l2"; oracle: "mlx-optiq"; claim: "bit-exact" }>
+  | Readonly<{ tier: "l3"; oracle: null; claim: "measured" }>;
+
+export type ModelLoader = "safetensors" | "colibri";
+export type ModelGraph =
+  | "gemma4"
+  | "minicpm5"
+  | "qwen3.5"
+  | "qwen3"
+  | "qwen3-moe"
+  | "diffusion-gemma"
+  | "glm5.2"
+  | "universal-dense";
+export type GenerationLoop = "autoregressive" | "diffusion";
+export type ModelSpecialization = "artifact" | "dedicated" | "generated" | "generic";
+
+export interface ModelExecutionComposition {
+  readonly loader: ModelLoader;
+  readonly graph: ModelGraph;
+  readonly loop: GenerationLoop;
+  readonly specialization: ModelSpecialization;
+  /** Engine-owned code registration. Model files never supply executable code.
+   * Omit to retain the existing graph implementation. */
+  readonly implementation?: string;
+}
+
+/** A profile declares construction only. Request methods such as MTP, KV
+ * schemes, adapters, grammar, and sampling are resolved independently and
+ * cannot be rewritten by profile selection. */
+export interface ModelProfile {
+  readonly id: string;
+  /** Exact external identity. Family profiles omit this field. */
+  readonly artifactFingerprint?: string;
+  /** Structural guard for an exact artifact declaration. */
+  readonly configFingerprint?: string;
+  readonly fidelity: FidelityTarget;
+  readonly requiredCapabilities: readonly EngineCapability[];
+  readonly execution: ModelExecutionComposition;
+}
+
+export interface ArtifactModelProfile extends ModelProfile {
+  readonly artifactFingerprint: string;
+  readonly configFingerprint: string;
+}
+
+export interface ModelArtifactIdentity {
+  readonly fingerprint: string | null;
+  readonly configFingerprint: string;
+}
+
+export interface ResolvedModelProfile {
+  readonly profile: ModelProfile;
+  readonly artifact: ModelArtifactIdentity;
+  readonly exactArtifact: boolean;
+}
+
+export interface ResolveModelProfileOptions {
+  /** Explicit identity for a non-Hugging-Face artifact. The caller owns its
+   * provenance; mlx-bun never mistakes a local path for a content identity. */
+  readonly artifactFingerprint?: string | null;
+  /** Additional exact declarations composed with the shipped profiles. */
+  readonly artifactProfiles?: readonly ArtifactModelProfile[];
+  readonly engineCapabilities?: readonly EngineCapability[];
+}
+
+const L1: FidelityTarget = Object.freeze({
+  tier: "l1", oracle: "mlx-lm", claim: "bit-exact",
+});
+const L2: FidelityTarget = Object.freeze({
+  tier: "l2", oracle: "mlx-optiq", claim: "bit-exact",
+});
+const L3: FidelityTarget = Object.freeze({
+  tier: "l3", oracle: null, claim: "measured",
+});
+
+function freezeProfile<T extends ModelProfile>(profile: T): T {
+  Object.freeze(profile.fidelity);
+  Object.freeze(profile.requiredCapabilities);
+  Object.freeze(profile.execution);
+  return Object.freeze(profile);
+}
+
+function snapshotProfile(profile: ModelProfile): ModelProfile {
+  return freezeProfile({
+    ...profile,
+    fidelity: { ...profile.fidelity },
+    requiredCapabilities: [...profile.requiredCapabilities],
+    execution: { ...profile.execution },
+  });
+}
+
+/** Exact artifact declarations already backed by this repository's parity or
+ * measured evidence. A new revision gets the family profile until it earns a
+ * declaration of its own. */
+export const BUILTIN_ARTIFACT_PROFILES: readonly ArtifactModelProfile[] = Object.freeze([
+  freezeProfile({
+    id: "qwen3.8-27b-optiq-4bit",
+    artifactFingerprint:
+      "hf:mlx-community/Qwen3.8-27B-OptiQ-4bit@b04599de95d7a9bfbd7f208d347c0f10d9432a42",
+    configFingerprint: "50117975cb405944",
+    fidelity: L1,
+    requiredCapabilities: Object.freeze([
+      "safetensors", "autoregressive", "qwen3.5-graph", "recurrent-state", "vision-sidecar",
+    ]),
+    execution: Object.freeze({
+      loader: "safetensors",
+      graph: "qwen3.5",
+      loop: "autoregressive",
+      specialization: "artifact",
+    }),
+  }),
+  freezeProfile({
+    id: "glm5.2-colibri-int4-int8-mtp",
+    artifactFingerprint:
+      "hf:mateogrgic/GLM-5.2-colibri-int4-with-int8-mtp@3cc8db99b1b13fc79325d987ba3c1c430766b3b8",
+    configFingerprint: "8b49d0941e18f1b9",
+    fidelity: L3,
+    requiredCapabilities: Object.freeze([
+      "colibri-container", "autoregressive", "glm5.2-graph", "streamed-experts",
+    ]),
+    execution: Object.freeze({
+      loader: "colibri",
+      graph: "glm5.2",
+      loop: "autoregressive",
+      specialization: "artifact",
+    }),
+  }),
+]);
+
+const FAMILY_PROFILES = {
+  gemma4Generated: freezeProfile({
+    id: "gemma4-generated",
+    fidelity: L1,
+    requiredCapabilities: ["safetensors", "autoregressive", "gemma4-graph", "generated-graph"],
+    execution: {
+      loader: "safetensors", graph: "gemma4", loop: "autoregressive", specialization: "generated",
+    },
+  }),
+  gemma4: freezeProfile({
+    id: "gemma4-dedicated",
+    fidelity: L1,
+    requiredCapabilities: ["safetensors", "autoregressive", "gemma4-graph"],
+    execution: {
+      loader: "safetensors", graph: "gemma4", loop: "autoregressive", specialization: "dedicated",
+    },
+  }),
+  minicpm5: freezeProfile({
+    id: "minicpm5-dedicated",
+    fidelity: L1,
+    requiredCapabilities: ["safetensors", "autoregressive", "minicpm5-graph"],
+    execution: {
+      loader: "safetensors", graph: "minicpm5", loop: "autoregressive", specialization: "dedicated",
+    },
+  }),
+  qwen35: freezeProfile({
+    id: "qwen3.5-dedicated",
+    fidelity: L1,
+    requiredCapabilities: ["safetensors", "autoregressive", "qwen3.5-graph", "recurrent-state"],
+    execution: {
+      loader: "safetensors", graph: "qwen3.5", loop: "autoregressive", specialization: "dedicated",
+    },
+  }),
+  qwen3: freezeProfile({
+    id: "qwen3-dedicated",
+    fidelity: L1,
+    requiredCapabilities: ["safetensors", "autoregressive", "qwen3-graph"],
+    execution: {
+      loader: "safetensors", graph: "qwen3", loop: "autoregressive", specialization: "dedicated",
+    },
+  }),
+  qwen3Moe: freezeProfile({
+    id: "qwen3-moe-dedicated",
+    fidelity: L1,
+    requiredCapabilities: ["safetensors", "autoregressive", "qwen3-moe-graph"],
+    execution: {
+      loader: "safetensors", graph: "qwen3-moe", loop: "autoregressive", specialization: "dedicated",
+    },
+  }),
+  diffusionGemma: freezeProfile({
+    id: "diffusion-gemma-dedicated",
+    fidelity: L2,
+    requiredCapabilities: ["safetensors", "diffusion", "diffusion-gemma-graph"],
+    execution: {
+      loader: "safetensors", graph: "diffusion-gemma", loop: "diffusion", specialization: "dedicated",
+    },
+  }),
+  glm52: freezeProfile({
+    id: "glm5.2-colibri",
+    fidelity: L3,
+    requiredCapabilities: ["colibri-container", "autoregressive", "glm5.2-graph", "streamed-experts"],
+    execution: {
+      loader: "colibri", graph: "glm5.2", loop: "autoregressive", specialization: "dedicated",
+    },
+  }),
+  universal: freezeProfile({
+    id: "universal-dense",
+    fidelity: L1,
+    requiredCapabilities: ["safetensors", "autoregressive", "universal-dense-graph"],
+    execution: {
+      loader: "safetensors", graph: "universal-dense", loop: "autoregressive", specialization: "generic",
+    },
+  }),
+} satisfies Record<string, ModelProfile>;
+
+/** Return a stable external fingerprint only when the directory is an exact HF
+ * snapshot revision. Relocation does not change the identity. Arbitrary local
+ * directories and mutable aliases intentionally return null. */
+export function externalArtifactFingerprint(modelDir: string): string | null {
+  const match = modelDir.replaceAll("\\", "/").match(
+    /(?:^|\/)models--([^/]+)\/snapshots\/([0-9a-f]{40,64})(?:\/|$)/i,
+  );
+  if (!match) return null;
+  const encodedRepo = match[1]!;
+  const separator = encodedRepo.indexOf("--");
+  if (separator <= 0 || separator === encodedRepo.length - 2) return null;
+  const repo = `${encodedRepo.slice(0, separator)}/${encodedRepo.slice(separator + 2)}`;
+  return `hf:${repo}@${match[2]!.toLowerCase()}`;
+}
+
+function familyProfile(config: ModelConfig, fingerprint: string): ModelProfile {
+  if (isGlm52Config(config)) return FAMILY_PROFILES.glm52;
+  if (isDiffusionGemmaConfig(config)) return FAMILY_PROFILES.diffusionGemma;
+  if (isMiniCPM5Config(config)) return FAMILY_PROFILES.minicpm5;
+  if (isQwen35Config(config)) return FAMILY_PROFILES.qwen35;
+  if (isQwen3MoeConfig(config)) return FAMILY_PROFILES.qwen3Moe;
+  if (isQwen3Config(config)) return FAMILY_PROFILES.qwen3;
+  if (config.modelType.startsWith("gemma4"))
+    return GENERATED.has(fingerprint) ? FAMILY_PROFILES.gemma4Generated : FAMILY_PROFILES.gemma4;
+  if (genericArgsFor(config)) return FAMILY_PROFILES.universal;
+
+  const arch = remapModelType(config.modelType);
+  throw new Error(
+    `unsupported model_type "${config.modelType}"` +
+    (arch !== config.modelType ? ` (mlx-lm remaps it to "${arch}")` : "") +
+    ` — targeted: gemma4*, diffusion_gemma, glm_moe_dsa, qwen3_5, qwen3, qwen3_moe, MiniCPM5;` +
+    ` generic (Tier-0): ${[...GENERIC_MODEL_TYPES].sort().join(", ")}`,
+  );
+}
+
+interface GraphMetadata {
+  readonly accepts: (config: ModelConfig) => boolean;
+  readonly capabilities: readonly EngineCapability[];
+}
+
+const GRAPH_METADATA: Readonly<Record<ModelGraph, GraphMetadata>> = Object.freeze({
+  "gemma4": {
+    accepts: (config) => config.modelType.startsWith("gemma4"),
+    capabilities: ["gemma4-graph"],
+  },
+  "minicpm5": { accepts: isMiniCPM5Config, capabilities: ["minicpm5-graph"] },
+  "qwen3.5": {
+    accepts: isQwen35Config,
+    capabilities: ["qwen3.5-graph", "recurrent-state"],
+  },
+  "qwen3": { accepts: isQwen3Config, capabilities: ["qwen3-graph"] },
+  "qwen3-moe": { accepts: isQwen3MoeConfig, capabilities: ["qwen3-moe-graph"] },
+  "diffusion-gemma": {
+    accepts: isDiffusionGemmaConfig,
+    capabilities: ["diffusion-gemma-graph"],
+  },
+  "glm5.2": {
+    accepts: isGlm52Config,
+    capabilities: ["glm5.2-graph", "streamed-experts"],
+  },
+  "universal-dense": {
+    accepts: (config) => {
+      try { return genericArgsFor(config) !== null; } catch { return false; }
+    },
+    capabilities: ["universal-dense-graph"],
+  },
+});
+
+function graphAccepts(profile: ModelProfile, config: ModelConfig): boolean {
+  return GRAPH_METADATA[profile.execution.graph].accepts(config);
+}
+
+function executionCapabilities(profile: ModelProfile): EngineCapability[] {
+  const required: EngineCapability[] = [profile.execution.loader === "colibri"
+    ? "colibri-container"
+    : "safetensors"];
+  if (profile.execution.loop === "diffusion") required.push("diffusion");
+  else required.push("autoregressive");
+  required.push(...GRAPH_METADATA[profile.execution.graph].capabilities);
+  if (profile.execution.specialization === "generated") required.push("generated-graph");
+  return required;
+}
+
+function validateProfile(profile: ModelProfile): void {
+  if (!profile.id) throw new Error("model profile id must not be empty");
+  if (profile.execution.implementation !== undefined &&
+      (!profile.execution.implementation.trim() || !profile.artifactFingerprint ||
+       profile.execution.specialization !== "artifact"))
+    throw new Error(`model profile ${profile.id} must bind a named implementation to an exact artifact`);
+  if ((profile.artifactFingerprint === undefined) !== (profile.configFingerprint === undefined))
+    throw new Error(
+      `model profile ${profile.id} must declare artifactFingerprint and configFingerprint together`,
+    );
+  const fidelityOk =
+    (profile.fidelity.tier === "l1" && profile.fidelity.oracle === "mlx-lm" &&
+      profile.fidelity.claim === "bit-exact") ||
+    (profile.fidelity.tier === "l2" && profile.fidelity.oracle === "mlx-optiq" &&
+      profile.fidelity.claim === "bit-exact") ||
+    (profile.fidelity.tier === "l3" && profile.fidelity.oracle === null &&
+      profile.fidelity.claim === "measured");
+  if (!fidelityOk) throw new Error(`model profile ${profile.id} has an invalid fidelity contract`);
+  const expectedLoader = profile.execution.graph === "glm5.2" ? "colibri" : "safetensors";
+  const expectedLoop = profile.execution.graph === "diffusion-gemma"
+    ? "diffusion"
+    : "autoregressive";
+  if (profile.execution.loader !== expectedLoader || profile.execution.loop !== expectedLoop)
+    throw new Error(
+      `model profile ${profile.id} has an invalid execution composition for ` +
+      `${profile.execution.graph}`,
+    );
+  const declared = new Set(profile.requiredCapabilities);
+  const underdeclared = executionCapabilities(profile).filter((capability) => !declared.has(capability));
+  if (underdeclared.length)
+    throw new Error(
+      `model profile ${profile.id} does not declare execution capabilities: ` +
+      `${underdeclared.join(", ")}`,
+    );
+}
+
+export function resolveModelProfile(
+  config: ModelConfig,
+  options: ResolveModelProfileOptions = {},
+): ResolvedModelProfile {
+  const fingerprint = configFingerprint(config);
+  const artifactFingerprint = options.artifactFingerprint === undefined
+    ? externalArtifactFingerprint(config.modelDir)
+    : options.artifactFingerprint;
+  const artifactProfiles = options.artifactProfiles
+    ? [...BUILTIN_ARTIFACT_PROFILES, ...options.artifactProfiles]
+    : BUILTIN_ARTIFACT_PROFILES;
+  const capabilities = new Set(options.engineCapabilities ?? ENGINE_CAPABILITIES);
+
+  let profile: ModelProfile | undefined;
+  let exactArtifact = false;
+  if (artifactFingerprint) {
+    const matches = artifactProfiles.filter((entry) => {
+      validateProfile(entry);
+      return entry.artifactFingerprint === artifactFingerprint;
+    });
+    if (matches.length > 1)
+      throw new Error(`multiple model profiles declare artifact ${artifactFingerprint}`);
+    profile = matches[0];
+    if (profile) {
+      exactArtifact = true;
+      if (profile.configFingerprint !== fingerprint)
+        throw new Error(
+          `model profile ${profile.id} matched artifact ${artifactFingerprint}, but config fingerprint ` +
+          `${fingerprint} != declared ${profile.configFingerprint}; refusing to fall back`,
+        );
+      if (!graphAccepts(profile, config))
+        throw new Error(
+          `model profile ${profile.id} selects ${profile.execution.graph}, which is incompatible with ` +
+          `${config.modelType}; refusing to fall back`,
+        );
+    }
+  }
+  profile ??= familyProfile(config, fingerprint);
+  validateProfile(profile);
+  profile = snapshotProfile(profile);
+
+  const missing = profile.requiredCapabilities.filter((capability) => !capabilities.has(capability));
+  if (missing.length)
+    throw new Error(
+      `model profile ${profile.id} requires missing engine capabilities: ${missing.join(", ")}; ` +
+      `refusing to fall back`,
+    );
+
+  return Object.freeze({
+    profile,
+    artifact: Object.freeze({ fingerprint: artifactFingerprint, configFingerprint: fingerprint }),
+    exactArtifact,
+  });
+}
+
+/** Guard the factory seam when a caller passes a previously resolved profile. */
+export function assertResolvedModelProfile(
+  config: ModelConfig,
+  resolved: ResolvedModelProfile,
+): void {
+  validateProfile(resolved.profile);
+  const fingerprint = configFingerprint(config);
+  if (resolved.artifact.configFingerprint !== fingerprint)
+    throw new Error(
+      `model profile ${resolved.profile.id} was resolved for config ` +
+      `${resolved.artifact.configFingerprint}, not ${fingerprint}`,
+    );
+  if (!graphAccepts(resolved.profile, config))
+    throw new Error(
+      `model profile ${resolved.profile.id} selects ${resolved.profile.execution.graph}, ` +
+      `which is incompatible with ${config.modelType}`,
+    );
+  if (resolved.exactArtifact !== (resolved.profile.artifactFingerprint !== undefined) ||
+      (resolved.exactArtifact &&
+       (resolved.profile.artifactFingerprint !== resolved.artifact.fingerprint ||
+        resolved.profile.configFingerprint !== fingerprint)))
+    throw new Error(
+      `model profile ${resolved.profile.id} does not match artifact ` +
+      `${resolved.artifact.fingerprint ?? "<unidentified>"}`,
+    );
+}

package/src/model/qwen-append.ts

@@ -0,0 +1,12 @@
+// MLX 0.32.2, mlx/backend/metal/scaled_dot_product_attention.cpp (MIT):
+// https://github.com/ml-explore/mlx/blob/v0.32.2/mlx/backend/metal/scaled_dot_product_attention.cpp
+// For the qualified Qwen GQA-6/D256 graph on applegpu_g16s, one-pass/vector
+// selection and two-pass block counts change after these inclusive KV lengths.
+// A committed span must stay within one arithmetic regime to match M=1.
+const ATTENTION_BOUNDARIES = [1023, 1024, 8192, 32768, 65536];
+
+export function qwenAppendChunkSize(cacheTokens: number): number {
+  for (const end of ATTENTION_BOUNDARIES)
+    if (cacheTokens < end) return Math.min(4, end - cacheTokens);
+  return 4;
+}

package/src/model/qwen-conv.ts

@@ -0,0 +1,52 @@
+// Experimental Qwen GDN convolution. Its two outputs are evaluated together:
+// consuming the activation also materializes the small, independently owned
+// state tail, without keeping a concatenated prefill buffer alive for it.
+// Arithmetic follows MLX 0.31.2 depthwise_conv_1d and unary sigmoid (MIT,
+// Copyright Apple Inc.), preserving the bf16 intermediate rounding.
+import type { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+
+/** Borrow all inputs; return owned activation and an independent state tail. */
+export type QwenConvolution = (x: MlxArray, state: MlxArray, weight: MlxArray) => [MlxArray, MlxArray];
+
+let kernel: MetalKernel | null = null;
+export const fusedQwenConvolution: QwenConvolution = (x, state, weight) => {
+  const [B, M, D] = x.shape as [number, number, number];
+  const K = weight.shape[1]!;
+  if (x.shape.length !== 3 || weight.shape.length !== 3 || weight.shape[0] !== D ||
+      weight.shape[2] !== 1 || K !== 4 || state.shape.join() !== [B, K - 1, D].join() ||
+      ![Dtype.bfloat16, Dtype.float32].includes(x.dtype) || state.dtype !== x.dtype || weight.dtype !== x.dtype)
+    throw new Error("fused Qwen convolution requires matching bf16/f32 inputs, [B,M,D], [B,3,D], [D,4,1]");
+  kernel ??= new MetalKernel({ name: "qwen_causal_conv_silu_tail",
+    inputNames: ["state", "x", "w"], outputNames: ["out", "tail"],
+    source: String.raw`
+      const uint c = thread_position_in_grid.x;
+      const uint t = thread_position_in_grid.y;
+      const uint b = thread_position_in_grid.z;
+      if (c >= D) return;
+      if (t < M) {
+        float acc = 0.0f;
+        for (uint i = 0; i < K; ++i) {
+          const uint row = t + i;
+          const T value = row < K-1 ? state[(b*(K-1)+row)*D+c] : x[(b*M+row-(K-1))*D+c];
+          acc += float(value) * float(w[c*K+i]);
+        }
+        const T conv = T(acc);
+        // Preserve the pinned MLX sigmoid's typed operations and bf16
+        // intermediate rounding. Promoting this expression to f32 differs.
+        const auto y = 1 / (1 + metal::exp(metal::abs(conv)));
+        const T sig = conv < 0 ? y : 1-y;
+        out[(b*M+t)*D+c] = conv * sig;
+      }
+      if (t < K-1) {
+        const uint row = M + t;
+        tail[(b*(K-1)+t)*D+c] = row < K-1 ? state[(b*(K-1)+row)*D+c] : x[(b*M+row-(K-1))*D+c];
+      }
+    ` });
+  return kernel.apply([state, x, weight], {
+    outputs: [{ shape: [B, M, D], dtype: x.dtype }, { shape: [B, K - 1, D], dtype: x.dtype }],
+    grid: [D, Math.max(M, K - 1), B], threadGroup: [256, 1, 1],
+    templateDtypes: { T: x.dtype }, templateInts: { D, M, K },
+  }) as [MlxArray, MlxArray];
+};

package/src/model/qwen3-delta.ts

@@ -194,10 +194,37 @@ export function gatedDeltaUpdate(
  *  and writes these slots directly; `advance(N)` only tracks token count for
  *  B=1 single-stream (lengths / left_padding are batched-decode concerns).
  *  Not a KVCache, so maybeQuantizeKv skips it. */
+/** One armed speculative verify round on an SSMCache: the pre-round state
+ *  snapshot (retained refs — MLX arrays are immutable, so this is free) plus
+ *  the layer's recorded position-local kernel inputs, which make a partial
+ *  reject bit-exactly replayable: state after the first `keep` window tokens
+ *  is the same arithmetic prefix whether or not the rejected tail was ever
+ *  processed. The replay is installed by the OWNING layer (it needs the conv
+ *  weight / A_log / dt_bias); the cache only stores and frees. */
+export interface SsmSpecRound {
+  /** True from specRoundBegin() until the layer's forward records. */
+  armed: boolean;
+  prevConv: MlxArray | null;
+  prevRecurrent: MlxArray | null;
+  prevOffset: number;
+  /** Pre-conv in_proj output [B,S,convDim] — position-local. */
+  qkv: MlxArray | null;
+  /** Gate/beta projections [B,S,Hv] — position-local. */
+  a: MlxArray | null;
+  b: MlxArray | null;
+  /** Window length S of the recorded verify forward. */
+  S: number;
+  /** Layer-bound prefix replay: called AFTER the snapshot is restored onto
+   *  the cache; advances conv/recurrent/offset by `keep` tokens. */
+  replay: ((cache: SSMCache, keep: number) => void) | null;
+}
+
 export class SSMCache implements Cache {
   conv: MlxArray | null = null;
   recurrent: MlxArray | null = null;
   offset = 0;
+  /** Live speculative verify round (serial spec lane only; null otherwise). */
+  specRound: SsmSpecRound | null = null;
   /** Per-row token coverage (batch lane only; null on serial B=1 caches,
    *  where `offset` IS the row's count). mlx-lm's ArraysCache tracks no
    *  offset at all (cache.py:594-722) — per-sequence coverage is the
@@ -207,6 +234,10 @@ export class SSMCache implements Cache {
    *  lockstep by advance(), row-filtered by filter(). */
   offsets: number[] | null = null;
 
+  signature(): string { return "ssm"; }
+  bytesPerToken(): number { return 0; }
+  get batchSize(): number { return this.offsets?.length ?? 1; }
+
   /** Row `i`'s own token coverage (per-row when batched, `offset` serial). */
   rowOffset(i: number): number {
     return this.offsets ? this.offsets[i]! : this.offset;
@@ -241,14 +272,77 @@ export class SSMCache implements Cache {
   }
 
   isTrimmable(): boolean {
-    return false; // recurrent state can't be rolled back per token
+    return false; // recurrent state can't be trimmed per token — see specRound*
   }
 
   trim(_n: number): void {
     throw new Error("SSMCache is not trimmable");
   }
 
+  // --- speculative verify-round support (serial spec lane, B=1) ------------
+  // The serve loop arms a round before the verify forward and resolves it
+  // after the accept walk. The layer's forward, seeing an armed round, hands
+  // its replaced state slots to the round instead of disposing them, records
+  // qkv/a/b, and installs the replay. Cache/interface contract:
+  // rollback(keep) restores the snapshot then replays `keep` window tokens
+  // bit-exactly; commit frees the snapshot + recordings.
+
+  specRoundBegin(): void {
+    this.#dropSpecRound(); // defensive: a mid-round throw left a stale round
+    this.specRound = {
+      armed: true,
+      prevConv: null,
+      prevRecurrent: null,
+      prevOffset: this.offset,
+      qkv: null,
+      a: null,
+      b: null,
+      S: 0,
+      replay: null,
+    };
+  }
+
+  specRoundCommit(): void {
+    this.#dropSpecRound();
+  }
+
+  specRoundRollback(keep: number): void {
+    const r = this.specRound;
+    if (!r) throw new Error("SSMCache.specRoundRollback without an armed round");
+    if (r.armed)
+      throw new Error("SSMCache.specRoundRollback before the verify forward recorded");
+    if (keep < 0 || keep > r.S)
+      throw new Error(`SSMCache.specRoundRollback keep=${keep} outside window S=${r.S}`);
+    // Restore the pre-round snapshot (ownership moves back to the cache).
+    this.conv?.dispose();
+    this.recurrent?.dispose();
+    this.conv = r.prevConv;
+    this.recurrent = r.prevRecurrent;
+    r.prevConv = null;
+    r.prevRecurrent = null;
+    this.offset = r.prevOffset;
+    if (this.offsets)
+      throw new Error("SSMCache.specRoundRollback on a batched cache (serial lane only)");
+    if (keep > 0) {
+      if (!r.replay) throw new Error("SSMCache.specRoundRollback with no recorded replay");
+      r.replay(this, keep); // advances conv/recurrent/offset by `keep`
+    }
+    this.#dropSpecRound();
+  }
+
+  #dropSpecRound(): void {
+    const r = this.specRound;
+    if (!r) return;
+    r.prevConv?.dispose();
+    r.prevRecurrent?.dispose();
+    r.qkv?.dispose();
+    r.a?.dispose();
+    r.b?.dispose();
+    this.specRound = null;
+  }
+
   dispose(): void {
+    this.#dropSpecRound();
     this.conv?.dispose();
     this.recurrent?.dispose();
     this.conv = null;
@@ -304,6 +398,8 @@ export class SSMCache implements Cache {
     if (this.offsets) this.offsets = keep.map((i) => this.offsets![i]!);
   }
 
+  filterRows(keep: readonly number[]): void { this.filter([...keep]); }
+
   /** Row `i` as a fresh SERIAL cache — port of mlx-lm ArraysCache.extract
    *  (cache.py:673-676: `cache.cache = [c[idx : idx + 1] for c in self.cache]`,
    *  a B-axis slice of every state slot). Ours are OWNED contiguous copies
@@ -319,7 +415,7 @@ export class SSMCache implements Cache {
       const hi = [...a.shape];
       hi[0] = i + 1;
       const view = a.slice(lo, hi);
-      const own = ops.contiguous(view);
+      const own = ops.copyOf(view); // TRUE copy: contiguous(view) is a no-op VIEW when already contiguous — pins the source buffer (2026-08-20 DeltaNet conv leak class)
       view.dispose();
       return own;
     };