mlx-bun 0.2.0 → 0.4.0

package/src/model/implementation.ts

@@ -0,0 +1,69 @@
+import type { ModelConfig } from "../config";
+import {
+  assertResolvedModelProfile,
+  type GenerationLoop,
+  type ModelGraph,
+  type ModelLoader,
+  type ResolvedModelProfile,
+} from "./profile";
+
+/** Engine code implements this port. Source and Model remain backend-specific;
+ * selection introduces no tensor conversion or dispatch in the token loop.
+ * create borrows source; a failed construction must release its own resources. */
+export interface ModelImplementation<Source, Model> {
+  readonly id: string;
+  readonly graph: ModelGraph;
+  readonly loader: ModelLoader;
+  readonly loop: GenerationLoop;
+  create(source: Source, config: ModelConfig, profile: ResolvedModelProfile): Model;
+}
+
+export interface ModelImplementationProvider<Source, Model> {
+  /** Validate identity and composition, then return one compatible binding.
+   * A declared implementation must never silently fall back to another. */
+  select(config: ModelConfig, profile: ResolvedModelProfile): ModelImplementation<Source, Model>;
+}
+
+/** Immutable composition, with identity matching delegated to model profiles.
+ * Exact declarations are requirements: absent/incompatible code is an error,
+ * never an invitation to retry with a different numerical implementation. */
+export class ModelImplementationRegistry<Source, Model> implements ModelImplementationProvider<Source, Model> {
+  private readonly entries: ReadonlyMap<string, ModelImplementation<Source, Model>>;
+
+  constructor(implementations: readonly ModelImplementation<Source, Model>[]) {
+    const entries = new Map<string, ModelImplementation<Source, Model>>();
+    for (const implementation of implementations) {
+      if (!implementation.id.trim()) throw new Error("model implementation id must not be empty");
+      if (entries.has(implementation.id))
+        throw new Error(`duplicate model implementation ${implementation.id}`);
+      entries.set(implementation.id, Object.freeze({
+        id: implementation.id,
+        graph: implementation.graph,
+        loader: implementation.loader,
+        loop: implementation.loop,
+        create: implementation.create.bind(implementation),
+      }));
+    }
+    this.entries = entries;
+  }
+
+  /** Return a new composition; active models retain their existing binding. */
+  with(...implementations: readonly ModelImplementation<Source, Model>[]): ModelImplementationRegistry<Source, Model> {
+    return new ModelImplementationRegistry([...this.entries.values(), ...implementations]);
+  }
+
+  select(config: ModelConfig, profile: ResolvedModelProfile): ModelImplementation<Source, Model> {
+    assertResolvedModelProfile(config, profile);
+    const execution = profile.profile.execution;
+    // These default IDs preserve the existing dedicated/generated/family choice.
+    const id = execution.implementation ?? (execution.specialization === "generated"
+      ? `${execution.graph}-generated` : execution.graph);
+    const implementation = this.entries.get(id);
+    if (!implementation)
+      throw new Error(`model profile ${profile.profile.id} requires unavailable implementation ${id}; refusing to fall back`);
+    if (implementation.graph !== execution.graph || implementation.loader !== execution.loader ||
+        implementation.loop !== execution.loop)
+      throw new Error(`model implementation ${id} is incompatible with profile ${profile.profile.id}; refusing to fall back`);
+    return implementation;
+  }
+}

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

@@ -0,0 +1,67 @@
+import type { MlxArray } from "../mlx/array";
+import * as ops from "../mlx/ops";
+import type { Cache, KvAttentionView, KvDonorAttention } from "./gemma4-base";
+import { quantizedSdpa, disposeTriple } from "./gemma4-base";
+
+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(); },
+  };
+}
+
+export function quantizedDonorAttention(keys: ops.QuantizedTensor, values: ops.QuantizedTensor,
+  validity: Pick<KvDonorAttention, "offsets" | "starts" | "ends">,
+  groupSize: number, bits: number): KvDonorAttention {
+  return {
+    width: keys.packed.shape[2]!, dtype: keys.scales.dtype, ...validity,
+    attend: (q, scale, mask) => quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
+    dispose() { disposeTriple(keys); disposeTriple(values); },
+  };
+}
+
+export function captureKvAttention(cache: Cache, k: MlxArray, v: MlxArray): KvAttentionView {
+    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) => quantizedSdpa(q, keys, values, scale, mask, groupSize, bits),
+        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/model/minicpm5.ts

@@ -14,7 +14,6 @@ import {
   KVCache,
   LoraState,
   QuantizedKVCache,
-  TurboQuantKVCache,
   QuantizedEmbedding,
   QuantizedLinear,
   quantizedSdpa,
@@ -30,7 +29,7 @@ import { runtimeValue } from "../runtime-config";
  *    `@partial(mx.compile, shapeless=True) def swiglu(gate, x): return nn.silu(gate) * x`
  *  mx.compile fuses sigmoid + mul + mul into ONE kernel (mlx-lm's decode
  *  `CV2ISigmoid…Multiply`) instead of our three separate dispatches. Traced
- *  once, replayed thereafter. Used on the decode (M=1) path only. */
+ *  once, replayed thereafter for decode and prefill. */
 let _swigluClosure: CompiledFunction | null = null;
 export function compiledSwiglu(gate: MlxArray, up: MlxArray): MlxArray {
   if (!_swigluClosure) {
@@ -143,23 +142,28 @@ export class LlamaAttention {
     q = disposing(q, ropeStep(q));
     k = disposing(k, ropeStep(k));
     let attn: MlxArray;
-    if (cache instanceof QuantizedKVCache) {
-      const [keys, values] = cache.updateAndFetchQuantized(k, v);
+    if (cache.attentionState) {
+      const view = cache.attentionState.appendAndFetch(k, v);
+      k.dispose(); v.dispose();
+      try { attn = view.attend(q, this.scale, mask); } finally { view.dispose(); }
+    } else if (cache.quantizedAttention) {
+      const quantized = cache.quantizedAttention;
+      const [keys, values] = quantized.updateAndFetchQuantized(k, v);
       k.dispose();
       v.dispose();
-      attn = quantizedSdpa(q, keys, values, this.scale, mask, cache.groupSize, cache.bits);
+      attn = quantizedSdpa(q, keys, values, this.scale, mask, quantized.groupSize, quantized.bits);
       disposeTriple(keys);
       disposeTriple(values);
-    } else if (cache instanceof TurboQuantKVCache) {
+    } else if (cache.rotatedValueAttention) {
       // Deferred-V: sdpa runs on rotated values, then one InvFWHT on the
       // output (linear in V) — see TurboQuantKVCache.updateAndFetchDeferredV.
-      const [keys, values] = cache.updateAndFetchDeferredV(k, v);
+      const [keys, values] = cache.rotatedValueAttention.updateAndFetchDeferredV(k, v);
       k.dispose();
       v.dispose();
       const rotated = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr);
       keys.dispose();
       values.dispose();
-      attn = tqUnrotateValues(rotated);
+      attn = cache.rotatedValueAttention.captureValueTransform?.()(rotated) ?? tqUnrotateValues(rotated);
       rotated.dispose();
     } else {
       const [keys, values] = cache.updateAndFetch(k, v);
@@ -200,15 +204,13 @@ export class LlamaMLP {
 
   /** MLP: returns mlp(x); the caller adds the residual when given. */
   forward(x: MlxArray, residual?: MlxArray): MlxArray {
-    const H = x.shape[x.shape.length - 1]!;
-    const M = x.shape.reduce((a, b) => a * b, 1) / H;
     const gate = this.gate.forward(x);
     const up = this.up.forward(x);
-    // Decode (M=1): fuse silu·up into one kernel (mlx_lm activations.py swiglu).
-    // Plain sigmoid+mul+mul otherwise (M>1 prefill/training, or inside a compile
-    // trace) — identical math either way. MLX_BUN_COMPILED_SWIGLU=0 disables.
+    // Match mlx_lm activations.py at every sequence length. Separate kernels
+    // differ numerically on MLX 0.32.2 at larger prefill shapes. An enclosing
+    // compiled trace owns fusion itself; the explicit opt-out stays available.
     let hidden: MlxArray;
-    if (M === 1 && !isCompiledTrace() && runtimeValue("MLX_BUN_COMPILED_SWIGLU") !== "0") {
+    if (!isCompiledTrace() && runtimeValue("MLX_BUN_COMPILED_SWIGLU") !== "0") {
       hidden = compiledSwiglu(gate, up);
       gate.dispose();
       up.dispose();
@@ -321,7 +323,7 @@ export class MiniCPM5Model {
    *  NEVER disposed — the caller owns it (for segmented backward `h` is a leaf
    *  whose lifetime is managed by the autograd closure); intra-range
    *  intermediates ARE disposed. Building block for `forwardLayers` and the
-   *  segmented-backward training path (docs/design/segmented-backward-training.md).
+   *  segmented-backward training path (docs/design/orpo-training.md).
    *  MiniCPM5 is plain full-attention with no KV-sharing, so each layer reads
    *  its own stateless `cache[i]` and a single causal mask. */
   runLayerRange(h: MlxArray, aIdx: number, bIdx: number, cache: Cache[]): MlxArray {

package/src/model/profile.ts

@@ -58,6 +58,9 @@ export interface ModelExecutionComposition {
   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
@@ -323,6 +326,10 @@ function executionCapabilities(profile: ModelProfile): EngineCapability[] {
 
 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`,
@@ -413,6 +420,7 @@ export function assertResolvedModelProfile(
   config: ModelConfig,
   resolved: ResolvedModelProfile,
 ): void {
+  validateProfile(resolved.profile);
   const fingerprint = configFingerprint(config);
   if (resolved.artifact.configFingerprint !== fingerprint)
     throw new Error(
@@ -424,8 +432,10 @@ export function assertResolvedModelProfile(
       `model profile ${resolved.profile.id} selects ${resolved.profile.execution.graph}, ` +
       `which is incompatible with ${config.modelType}`,
     );
-  if (resolved.exactArtifact &&
-      resolved.profile.artifactFingerprint !== resolved.artifact.fingerprint)
+  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-state.ts

@@ -0,0 +1,74 @@
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import { MetalKernel } from "../mlx/metal-kernel";
+import * as ops from "../mlx/ops";
+import { computeG } from "./qwen3-delta";
+
+// The same recurrent arithmetic as gated_delta_step, with no query/output
+// calculation. Each row replays only its retained prefix of the verify window.
+const SOURCE = String.raw`
+  const int T = Tin[0];
+  auto n = thread_position_in_grid.z;
+  auto b_idx = n / Hv;
+  auto hv_idx = n % Hv;
+  auto hk_idx = hv_idx / (Hv / Hk);
+  constexpr int n_per_t = Dk / 32;
+  auto dk_idx = thread_position_in_threadgroup.x;
+  auto dv_idx = thread_position_in_grid.y;
+  auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+  auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+  auto g_ = g + b_idx * T * Hv;
+  auto beta_ = beta + b_idx * T * Hv;
+  auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+  auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+  float state[n_per_t];
+  for (int i = 0; i < n_per_t; ++i)
+    state[i] = static_cast<float>(i_state[n_per_t * dk_idx + i]);
+  const int keep = lengths[b_idx];
+  for (int t = 0; t < keep; ++t) {
+    float kv_mem = 0.0f;
+    for (int i = 0; i < n_per_t; ++i) {
+      auto s_idx = n_per_t * dk_idx + i;
+      state[i] = state[i] * g_[hv_idx];
+      kv_mem += state[i] * k_[s_idx];
+    }
+    kv_mem = simd_sum(kv_mem);
+    auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+    for (int i = 0; i < n_per_t; ++i) {
+      auto s_idx = n_per_t * dk_idx + i;
+      state[i] = state[i] + k_[s_idx] * delta;
+    }
+    k_ += Hk * Dk;
+    v_ += Hv * Dv;
+    g_ += Hv;
+    beta_ += Hv;
+  }
+  for (int i = 0; i < n_per_t; ++i)
+    o_state[n_per_t * dk_idx + i] = state[i];
+`;
+
+let kernel: MetalKernel | undefined;
+
+/** State-only prefix replay. Inputs keep their original [B,T,...] strides;
+ * lengths are request-owned accepted counts, never scheduling decisions. */
+export function gatedDeltaState(
+  k: MlxArray, v: MlxArray, a: MlxArray, b: MlxArray,
+  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null,
+  lengths: readonly number[],
+): MlxArray {
+  const [B, T, Hk, Dk] = k.shape as [number, number, number, number];
+  const [, , Hv, Dv] = v.shape as [number, number, number, number];
+  using beta = ops.sigmoid(b);
+  using g = computeG(aLog, a, dtBias);
+  using zero = state ? null : ops.zeros([B, Hv, Dv, Dk], Dtype.float32);
+  using t = ops.fromInt32([T], [1]);
+  using kept = ops.fromInt32([...lengths], [B]);
+  kernel ??= new MetalKernel({ name: "gated_delta_prefix_state",
+    inputNames: ["k", "v", "g", "beta", "state_in", "Tin", "lengths"],
+    outputNames: ["state_out"], source: SOURCE, ensureRowContiguous: true });
+  return kernel.apply([k, v, g, beta, state ?? zero!, t, kept], {
+    outputs: [{ shape: [B, Hv, Dv, Dk], dtype: Dtype.float32 }],
+    grid: [32, Dv, B * Hv], threadGroup: [32, 4, 1],
+    templateInts: { Dk, Dv, Hk, Hv },
+  })[0]!;
+}

package/src/model/qwen3-delta.ts

@@ -5,8 +5,8 @@
 // default (use_kernel = not training), and its float accumulation / simd_sum
 // reduction order differ from the pure-ops fallback — so BIT-EXACT parity with
 // mlx-lm requires the SAME kernel, dispatched with the SAME grid/threadgroup.
-// We port the non-vectorized, non-masked variant (g.ndim == 3, mask is None),
-// which is the B=1 single-stream path (ssm_mask is None at batch 1).
+// We port the non-vectorized variant (g.ndim == 3). Padded prefill selects
+// the oracle masked specialization; unpadded execution retains its kernel.
 //
 // Numerics (must match the reference dtypes exactly — mlx infers the kernel's
 // pointer element types from the input arrays):
@@ -20,10 +20,12 @@
 // GQA is handled inside the kernel (hk_idx = hv_idx / (Hv/Hk)); q/k stay at Hk.
 
 import { MlxArray } from "../mlx/array";
+import { materializeCopy } from "../mlx/materialize";
 import { Dtype } from "../mlx/ffi";
 import { MetalKernel } from "../mlx/metal-kernel";
 import { CompiledFunction } from "../mlx/compile";
 import * as ops from "../mlx/ops";
+import type { SsmPrefillPadding } from "./ssm-prefill-padding";
 import type { Cache, Mask } from "./gemma4-base";
 
 // Verbatim body of mlx-lm's gated_delta_step (has_mask=False, vectorized=False),
@@ -104,7 +106,15 @@ const SOURCE = String.raw`
 `;
 
 let kernel: MetalKernel | null = null;
-function getKernel(): MetalKernel {
+let maskedKernel: MetalKernel | null = null;
+function getKernel(masked = false): MetalKernel {
+  if (masked) return maskedKernel ??= new MetalKernel({
+    name: "gated_delta_step_mask",
+    inputNames: ["q", "k", "v", "g", "beta", "state_in", "Tin", "mask"],
+    outputNames: ["y", "state_out"],
+    source: SOURCE.replace("if (true)", "if (mask[b_idx * T + t])"),
+    ensureRowContiguous: true,
+  });
   if (!kernel)
     kernel = new MetalKernel({
       name: "gated_delta_step",
@@ -154,7 +164,7 @@ export function computeG(aLog: MlxArray, a: MlxArray, dtBias: MlxArray): MlxArra
  *   state: [B, Hv, Dv, Dk] f32 (or null → zeros) */
 export function gatedDeltaUpdate(
   q: MlxArray, k: MlxArray, v: MlxArray, a: MlxArray, b: MlxArray,
-  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null,
+  aLog: MlxArray, dtBias: MlxArray, state: MlxArray | null, mask: MlxArray | null = null,
 ): [MlxArray, MlxArray] {
   const [B, , Hk, Dk] = q.shape as [number, number, number, number];
   const [, , Hv, Dv] = v.shape as [number, number, number, number];
@@ -171,7 +181,9 @@ export function gatedDeltaUpdate(
 
   const T = q.shape[1]!;
   const tArr = MlxArray.fromInt32(new Int32Array([T]), [1]);
-  const [y, stateOut] = getKernel().apply([q, k, v, g, beta, stateIn, tArr], {
+  const inputs = [q, k, v, g, beta, stateIn, tArr];
+  if (mask) inputs.push(mask);
+  const [y, stateOut] = getKernel(mask !== null).apply(inputs, {
     outputs: [
       { shape: [B, T, Hv, Dv], dtype: q.dtype },
       { shape: [B, Hv, Dv, Dk], dtype: Dtype.float32 },
@@ -207,6 +219,7 @@ export interface SsmSpecRound {
   prevConv: MlxArray | null;
   prevRecurrent: MlxArray | null;
   prevOffset: number;
+  prevOffsets: number[] | null;
   /** Pre-conv in_proj output [B,S,convDim] — position-local. */
   qkv: MlxArray | null;
   /** Gate/beta projections [B,S,Hv] — position-local. */
@@ -216,14 +229,15 @@ export interface SsmSpecRound {
   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;
+  replay: ((cache: SSMCache, keep: number | readonly 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). */
+  prefillPadding: SsmPrefillPadding | null = null;
+  /** Live speculative verify round; null outside a verify transaction. */
   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
@@ -248,15 +262,13 @@ export class SSMCache implements Cache {
     throw new Error("SSMCache has no KV updateAndFetch (gated-DeltaNet layer)");
   }
 
-  makeMask(_N: number, _windowSize: number | null): Mask {
-    // ssm_mask is None both single-stream and in the batch lane: rows
-    // solo-prefill UNPADDED at B=1 (state never sees pad tokens), and
-    // batched decode feeds one real token per row — so unlike mlx-lm's
-    // left-padded batch-prefill there is never a pad position to mask.
-    return { mode: "", arr: null };
+  makeMask(N: number, _windowSize: number | null): Mask {
+    const arr = this.prefillPadding?.makeMask(N) ?? null;
+    return { mode: arr ? "array" : "", arr };
   }
 
   advance(n: number): void {
+    if (this.prefillPadding) { this.advanceRows(this.prefillPadding.advance(n)); return; }
     this.offset += n;
     // Batched rows step together (one forward advances every row), so the
     // per-row counts move in lockstep; they DIFFER only in their merge-time
@@ -264,6 +276,13 @@ export class SSMCache implements Cache {
     if (this.offsets) for (let i = 0; i < this.offsets.length; i++) this.offsets[i]! += n;
   }
 
+  /** Methods may retain different verified prefixes for each active request. */
+  advanceRows(counts: readonly number[]): void {
+    if (!this.offsets) { this.offset += counts[0]!; return; }
+    for (let row = 0; row < this.offsets.length; row++) this.offsets[row]! += counts[row]!;
+    this.offset = Math.max(...this.offsets);
+  }
+
   state(): MlxArray[] {
     const out: MlxArray[] = [];
     if (this.conv) out.push(this.conv);
@@ -279,7 +298,7 @@ export class SSMCache implements Cache {
     throw new Error("SSMCache is not trimmable");
   }
 
-  // --- speculative verify-round support (serial spec lane, B=1) ------------
+  // --- speculative verify-round support -----------------------------------
   // 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
@@ -294,6 +313,7 @@ export class SSMCache implements Cache {
       prevConv: null,
       prevRecurrent: null,
       prevOffset: this.offset,
+      prevOffsets: this.offsets ? [...this.offsets] : null,
       qkv: null,
       a: null,
       b: null,
@@ -306,12 +326,12 @@ export class SSMCache implements Cache {
     this.#dropSpecRound();
   }
 
-  specRoundRollback(keep: number): void {
+  specRoundRollback(keep: number | readonly 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)
+    if (typeof keep === "number" ? keep < 0 || keep > r.S : keep.some(count => count < 0 || count > 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();
@@ -321,9 +341,8 @@ export class SSMCache implements Cache {
     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) {
+    this.offsets = r.prevOffsets;
+    if (typeof keep === "number" ? keep > 0 : keep.some(count => count > 0)) {
       if (!r.replay) throw new Error("SSMCache.specRoundRollback with no recorded replay");
       r.replay(this, keep); // advances conv/recurrent/offset by `keep`
     }
@@ -347,6 +366,7 @@ export class SSMCache implements Cache {
     this.recurrent?.dispose();
     this.conv = null;
     this.recurrent = null;
+    this.prefillPadding = null;
   }
 
   // --- batch-lane dynamic-B ops (BatchScheduler only) ----------------------
@@ -386,25 +406,30 @@ export class SSMCache implements Cache {
 
   /** Evict rows not in `keep` (ascending row indices), in place. */
   filter(keep: number[]): void {
-    if (!this.conv || !this.recurrent) return;
-    const idx = ops.fromInt32(keep, [keep.length]);
-    const conv = ops.takeAxis(this.conv, idx, 0);
-    const recurrent = ops.takeAxis(this.recurrent, idx, 0);
-    idx.dispose();
-    this.conv.dispose();
-    this.recurrent.dispose();
-    this.conv = conv;
-    this.recurrent = recurrent;
-    if (this.offsets) this.offsets = keep.map((i) => this.offsets![i]!);
+    if (this.conv && this.recurrent) {
+      const idx = ops.fromInt32(keep, [keep.length]);
+      const conv = ops.takeAxis(this.conv, idx, 0);
+      const recurrent = ops.takeAxis(this.recurrent, idx, 0);
+      idx.dispose();
+      this.conv.dispose();
+      this.recurrent.dispose();
+      this.conv = conv;
+      this.recurrent = recurrent;
+    }
+    this.prefillPadding?.filter(keep);
+    if (this.offsets) {
+      this.offsets = keep.map((i) => this.offsets![i]!);
+      this.offset = Math.max(0, ...this.offsets);
+    }
   }
 
   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
-   *  (the entry outlives the batch; the batched buffers must free when the
-   *  batch moves on — same rule as extractKVRow, batched-mask.ts). Offset =
+   *  a B-axis slice of every state slot). A proper subset gets compact storage
+   *  so retaining one row does not pin the other rows. A singleton retains an
+   *  independent handle over its complete allocation. Offset =
    *  the row's OWN coverage, not the shared max. */
   extractRow(i: number): SSMCache {
     if (!this.conv || !this.recurrent)
@@ -414,10 +439,8 @@ export class SSMCache implements Cache {
       lo[0] = i;
       const hi = [...a.shape];
       hi[0] = i + 1;
-      const view = a.slice(lo, hi);
-      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;
+      using view = a.slice(lo, hi);
+      return a.shape[0] === 1 ? ops.contiguous(view) : materializeCopy(view);
     };
     const out = new SSMCache();
     out.conv = cut(this.conv);