mlx-bun 0.0.10 → 0.0.12

package/src/model/minicpm5.ts

@@ -6,6 +6,7 @@ import type { Weights } from "../weights";
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import * as ops from "../mlx/ops";
+import { unrotateValues as tqUnrotateValues } from "../mlx/turboquant-ops";
 import {
   argmaxLastPosition,
   disposeTriple,
@@ -13,6 +14,7 @@ import {
   KVCache,
   LoraState,
   QuantizedKVCache,
+  TurboQuantKVCache,
   QuantizedEmbedding,
   QuantizedLinear,
   quantizedSdpa,
@@ -21,52 +23,27 @@ import {
   type Cache,
   type Mask,
 } from "./gemma4-base";
-import {
-  fusedSwiglu,
-  fusedSwigluEnabled,
-  fusedSwigluSupported,
-  swigluMMin,
-  swigluUnderCap,
-} from "./fused-swiglu-kernel";
-import { steelLinear } from "./steel-linear-kernel";
-import { fusedMlp } from "./fused-mlp-kernel";
-
-/** Break the graph at a boundary: contiguous → eval → fresh graph-free leaf — the
- *  exact primitive the segmented TRAINER uses so MLX never plans across boundaries
- *  and prior-segment activations are freed (copy of train/segmented.ts detachLeaf). */
-function detachLeaf(a: MlxArray): MlxArray {
-  const c = ops.contiguous(a);
-  c.eval();
-  const leaf = MlxArray.fromBytesCopy(c.rawBytes(), c.shape, c.dtype);
-  c.dispose();
-  a.dispose();
-  return leaf;
-}
+import { CompiledFunction } from "../mlx/compile";
 
-/** A quantized projection as OUR steel dispatch when the controlled path is live
- *  (no adapter, affine, clean tiling, M>1 so it's not a decode GEMV); otherwise
- *  the normal linear. x is [B,L,H] or [M,H]. */
-function ourLinear(lin: QuantizedLinear, x: MlxArray): MlxArray {
-  const H = x.shape[x.shape.length - 1]!;
-  const M = x.shape.reduce((a, b) => a * b, 1) / H;
-  if (
-    !fusedSwigluEnabled() || isCompiledTrace() || M < swigluMMin() ||
-    lin.adapters != null || lin.biases == null ||
-    (x.dtype !== Dtype.bfloat16 && x.dtype !== Dtype.float16) ||
-    H % 32 !== 0 || lin.spec.groupSize % 32 !== 0
-  ) {
-    return lin.forward(x);
+/** Verbatim port of mlx_lm/models/activations.py:
+ *    `@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. */
+let _swigluClosure: CompiledFunction | null = null;
+export function compiledSwiglu(gate: MlxArray, up: MlxArray): MlxArray {
+  if (!_swigluClosure) {
+    _swigluClosure = new CompiledFunction((inputs) => {
+      const g = inputs[0]!, u = inputs[1]!;              // nn.silu(gate) * x
+      const sig = ops.sigmoid(g);
+      const silu = ops.mul(g, sig); sig.dispose();
+      const out = ops.mul(silu, u); silu.dispose();
+      return [out];
+    });
   }
-  const x2 = x.shape.length === 2 ? x : ops.reshape(x, [M, H]);
-  const y = steelLinear(x2, lin.w, lin.scales, lin.biases, lin.spec);
-  if (x2 !== x) x2.dispose();
-  return x.shape.length === 2
-    ? y
-    : disposing(y, ops.reshape(y, [...x.shape.slice(0, -1), lin.w.shape[0]!]));
+  return _swigluClosure.apply([gate, up])[0]!;
 }
 
-/** How often the inference fused path materializes the residual stream (layers). */
-const EVAL_EVERY = Math.max(1, Number(process.env.MLX_BUN_SWIGLU_EVAL_EVERY) || 1);
 import { flashAttention, getTrainingAttn, flashSupported } from "./flash-attention";
 
 // Shared prompt-prefix plan (lever 7). When set, the attention ropes the
@@ -83,6 +60,11 @@ let _prefixPlan: PrefixPlan | null = null;
 export function setMiniCpmPrefixPlan(p: PrefixPlan | null): void {
   _prefixPlan = p;
 }
+/** True while a block-wise prefix-shared RoPE plan is active (ORPO forward).
+ *  The generated debranched fast path must fall back to the monolith then. */
+export function miniCpmPrefixPlanActive(): boolean {
+  return _prefixPlan != null;
+}
 
 /** Block-wise RoPE for the prefix-shared concat [prompt; chosen; rejected] along
  *  the sequence axis (axis 2 of [B,H,T,D]): prompt rotated at offset 0, each
@@ -108,7 +90,7 @@ function ropeBlocks(x: MlxArray, dims: number, base: number, plan: PrefixPlan):
   return out;
 }
 
-class LlamaAttention {
+export class LlamaAttention {
   readonly qProj: QuantizedLinear;
   readonly kProj: QuantizedLinear;
   readonly vProj: QuantizedLinear;
@@ -167,6 +149,17 @@ class LlamaAttention {
       attn = quantizedSdpa(q, keys, values, this.scale, mask, cache.groupSize, cache.bits);
       disposeTriple(keys);
       disposeTriple(values);
+    } else if (cache instanceof TurboQuantKVCache) {
+      // 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);
+      k.dispose();
+      v.dispose();
+      const rotated = ops.sdpa(q, keys, values, this.scale, mask.mode, mask.arr);
+      keys.dispose();
+      values.dispose();
+      attn = tqUnrotateValues(rotated);
+      rotated.dispose();
     } else {
       const [keys, values] = cache.updateAndFetch(k, v);
       k.dispose();
@@ -193,7 +186,7 @@ class LlamaAttention {
   }
 }
 
-class LlamaMLP {
+export class LlamaMLP {
   readonly gate: QuantizedLinear;
   readonly up: QuantizedLinear;
   readonly down: QuantizedLinear;
@@ -204,49 +197,29 @@ class LlamaMLP {
     this.down = QuantizedLinear.load(weights, `${prefix}.down_proj`, config);
   }
 
-  /** MLP. When `residual` (the layer input h) is given AND the fused path applies,
-   *  the kernel computes h + mlp(x) directly — the residual is ABSORBED, so nothing
-   *  stays live across the kernel for MLX's planner to alias (it just runs our one
-   *  dispatch). Otherwise returns mlp(x) and the caller adds the residual. */
+  /** 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;
-    if (
-      residual != null &&
-      fusedSwigluEnabled() && !isCompiledTrace() && M >= swigluMMin() && swigluUnderCap() &&
-      this.gate.adapters == null && this.up.adapters == null && this.down.adapters == null &&
-      this.down.biases != null && this.down.spec.groupSize % 32 === 0 &&
-      fusedSwigluSupported(x, this.gate.spec, this.up.spec, this.gate.biases, this.up.biases)
-    ) {
-      // The WHOLE MLP + residual in ONE kernel — gate+up+silu+down then +h, `hidden`
-      // never materialized, residual absorbed. MLX runs one dispatch; nothing for
-      // its planner to alias → chains across layers with no boundary.
-      const Hin = this.gate.scales.shape[1]! * this.gate.spec.groupSize;
-      const Ho = this.down.w.shape[0]!;
-      const x2 = x.shape.length === 2 ? x : ops.reshape(x, [M, Hin]);
-      const h2 = residual.shape.length === 2 ? residual : ops.reshape(residual, [M, Ho]);
-      const out2 = fusedMlp(
-        x2,
-        this.gate.w, this.gate.scales, this.gate.biases!, this.gate.spec,
-        this.up.w, this.up.scales, this.up.biases!, this.up.spec,
-        this.down.w, this.down.scales, this.down.biases!, this.down.spec,
-        h2,
-      );
-      if (x2 !== x) x2.dispose();
-      if (h2 !== residual) h2.dispose();
-      return x.shape.length === 2
-        ? out2
-        : disposing(out2, ops.reshape(out2, [...x.shape.slice(0, -1), Ho]));
-    }
     const gate = this.gate.forward(x);
     const up = this.up.forward(x);
-    const sig = ops.sigmoid(gate);
-    const silu = ops.mul(gate, sig);
-    gate.dispose();
-    sig.dispose();
-    const hidden = ops.mul(silu, up);
-    silu.dispose();
-    up.dispose();
+    // 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.
+    let hidden: MlxArray;
+    if (M === 1 && !isCompiledTrace() && process.env.MLX_BUN_COMPILED_SWIGLU !== "0") {
+      hidden = compiledSwiglu(gate, up);
+      gate.dispose();
+      up.dispose();
+    } else {
+      const sig = ops.sigmoid(gate);
+      const silu = ops.mul(gate, sig);
+      gate.dispose();
+      sig.dispose();
+      hidden = ops.mul(silu, up);
+      silu.dispose();
+      up.dispose();
+    }
     const m = this.down.forward(hidden);
     hidden.dispose();
     if (residual == null) return m;
@@ -254,7 +227,7 @@ class LlamaMLP {
   }
 }
 
-class LlamaLayer {
+export class LlamaLayer {
   readonly attn: LlamaAttention;
   readonly mlp: LlamaMLP;
   readonly inputNorm: RMSNorm;
@@ -337,15 +310,7 @@ export class MiniCPM5Model {
   }
 
   protected forwardLayers(h0: MlxArray, cache: Cache[]): MlxArray {
-    // When the fused SwiGLU path is live, WE own the eval schedule: materialize
-    // each layer's output so the custom kernel sits in a shallow per-layer graph
-    // rather than a 24-deep lazy graph whose buffer planner aliases the kernel's
-    // output (the deep-graph corruption). Training/segmented-backward call
-    // runLayerRange directly and never set this, so autograd is untouched.
-    const h = this.runLayerRange(
-      h0, 0, this.layers.length, cache,
-      fusedSwigluEnabled() && process.env.MLX_BUN_SWIGLU_NOEVAL !== "1",
-    );
+    const h = this.runLayerRange(h0, 0, this.layers.length, cache);
     h0.dispose(); // runLayerRange keeps its input alive; we own h0 here
     return disposing(h, this.finalNorm.forward(h));
   }
@@ -358,7 +323,7 @@ export class MiniCPM5Model {
    *  segmented-backward training path (docs/design/segmented-backward-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[], evalPerLayer = false): MlxArray {
+  runLayerRange(h: MlxArray, aIdx: number, bIdx: number, cache: Cache[]): MlxArray {
     if (aIdx >= bIdx) return h;
     const L = h.shape[1]!;
     const mask = cache[aIdx]!.makeMask(L, null);
@@ -367,15 +332,6 @@ export class MiniCPM5Model {
       const next = this.layers[i]!.forward(cur, mask, cache[i]!);
       if (i > aIdx) cur.dispose(); // keep the input boundary leaf; drop interiors
       cur = next;
-      // Per-layer boundary (segmented-trainer pattern): materialize each layer's
-      // output so MLX never plans across layers — our single-dispatch fused MLP is
-      // the only custom kernel it runs per layer. evalAll (GPU sync, no host copy)
-      // suffices for one kernel/boundary; MLX_BUN_SWIGLU_DETACH=1 forces the
-      // stronger graph-breaking detachLeaf copy.
-      if (evalPerLayer) {
-        if (process.env.MLX_BUN_SWIGLU_DETACH === "1") cur = detachLeaf(cur);
-        else ops.evalAll([cur]);
-      }
     }
     mask.arr?.dispose();
     return cur;

package/src/model/paged-kv.ts

@@ -0,0 +1,253 @@
+// Paged KV cache — vLLM-style block-pool storage behind the standard Cache
+// interface (docs/design/paged-kv-cache.md). OPTIONAL and default-off
+// (`--paged-kv` / GenerateOptions.pagedKv); v1 scope is serial batch=1
+// decode on Gemma4-family plain full-attention layers, bf16 only.
+//
+// Storage: K/V live in fixed-size per-layer pool tensors
+// [numBlocks, H_kv, blockSize, headDim]; a host-side block table maps the
+// sequence's logical block order to physical pool slots. Writes go through
+// ops.sliceUpdate into the tail block; reads gather the occupied blocks
+// back into ONE contiguous [1, H, S, D] pair (ops.takeAxis over the pool's
+// block axis + transpose + reshape) and hand that to the unchanged
+// ops.sdpa call site — identical bytes to what a plain KVCache would have
+// fetched, so the paged path is gated bit-exact, not KL-tolerated
+// (tests/paged-kv-parity.test.ts). The gather IS the cost: a full K/V copy
+// per step, pure bandwidth tax at batch=1 — v1 ships the abstraction for
+// the batched/CoW follow-ups, not a speed win (see the design doc's
+// Motivation, which says this honestly).
+//
+// Deliberately NOT a KVCache subclass (the TurboQuantKVCache reasoning):
+// every instanceof gate in the tree — CompiledDecode.supports, the batch
+// gateway's #modelCachesBatchable, generated forwards' #matches() — must
+// EXCLUDE paged caches and fall back to the monolith/uncompiled path.
+
+import { MlxArray } from "../mlx/array";
+import { Dtype } from "../mlx/ffi";
+import * as ops from "../mlx/ops";
+import { createCausalMask, type Cache, type Mask } from "./gemma4-base";
+
+/** Typed pool-exhaustion error: a generation outgrew its pool. Sizing from
+ *  prompt+maxTokens at construction makes this unreachable in practice;
+ *  it exists so an accounting bug is a clear throw, never silent
+ *  truncation or an allocator OOM (the byte-capped-not-count-capped
+ *  lesson). */
+export class PagedPoolExhausted extends Error {
+  constructor(numBlocks: number) {
+    super(`paged KV pool exhausted (${numBlocks} blocks) — capacity accounting bug`);
+    this.name = "PagedPoolExhausted";
+  }
+}
+
+/** Per-layer block arena. One K and one V pool tensor, mlx-owned
+ *  (ops.zeros — no host-pointer alignment/dtor hazards), shaped
+ *  [numBlocks, H, blockSize, D]. alloc()/free() manage physical slot
+ *  indices; the owning PagedKVCache writes/gathers through the pool. */
+export class BlockPool {
+  keys: MlxArray;
+  values: MlxArray;
+  readonly numBlocks: number;
+  readonly blockSize: number;
+  #free: number[];
+
+  constructor(opts: {
+    numBlocks: number;
+    blockSize: number;
+    numKvHeads: number;
+    headDim: number;
+    /** V head dim when it differs from K's (attention_k_eq_v models don't
+     *  reach v1's gemma4 scope, but the pool stays shape-honest). */
+    vHeadDim?: number;
+    dtype: Dtype;
+  }) {
+    this.numBlocks = opts.numBlocks;
+    this.blockSize = opts.blockSize;
+    this.keys = ops.zeros(
+      [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.headDim], opts.dtype);
+    this.values = ops.zeros(
+      [opts.numBlocks, opts.numKvHeads, opts.blockSize, opts.vHeadDim ?? opts.headDim],
+      opts.dtype);
+    // LIFO free list, low indices first — deterministic layout for tests.
+    this.#free = Array.from({ length: opts.numBlocks }, (_, i) => opts.numBlocks - 1 - i);
+  }
+
+  get freeBlocks(): number {
+    return this.#free.length;
+  }
+
+  alloc(): number {
+    const idx = this.#free.pop();
+    if (idx === undefined) throw new PagedPoolExhausted(this.numBlocks);
+    return idx;
+  }
+
+  free(physIdx: number): void {
+    this.#free.push(physIdx);
+  }
+
+  /** Write `k`/`v` ([1, H, l, D], l ≤ blockSize − within) into physical
+   *  block `physIdx` starting at row `within`. Functional slice_update:
+   *  the pool tensors are single-referenced here, so mlx donates the
+   *  buffer and the write is in-place in the steady state. */
+  writeBlock(physIdx: number, k: MlxArray, v: MlxArray, within: number): void {
+    const l = k.shape[2]!;
+    const [, H, , kD] = this.keys.shape as [number, number, number, number];
+    const vD = this.values.shape[3]!;
+    const k2 = ops.sliceUpdate(
+      this.keys, k, [physIdx, 0, within, 0], [physIdx + 1, H, within + l, kD]);
+    this.keys.dispose();
+    this.keys = k2;
+    const v2 = ops.sliceUpdate(
+      this.values, v, [physIdx, 0, within, 0], [physIdx + 1, H, within + l, vD]);
+    this.values.dispose();
+    this.values = v2;
+  }
+
+  /** Gather `blockTable`'s blocks in logical order into one contiguous
+   *  [1, H, nb·blockSize, D] pair (the fetch copy). Caller owns both. */
+  gather(blockTable: number[]): [MlxArray, MlxArray] {
+    const idx = MlxArray.fromInt32(Int32Array.from(blockTable), [blockTable.length]);
+    const pick = (pool: MlxArray): MlxArray => {
+      const took = ops.takeAxis(pool, idx, 0); // [nb, H, bs, D]
+      const [nb, H, bs, D] = took.shape as [number, number, number, number];
+      const t = ops.transposeAxes(took, [1, 0, 2, 3]); // [H, nb, bs, D] (view)
+      took.dispose();
+      const out = ops.reshape(t, [1, H, nb * bs, D]); // materializes seq order
+      t.dispose();
+      return out;
+    };
+    const k = pick(this.keys);
+    const v = pick(this.values);
+    idx.dispose();
+    return [k, v];
+  }
+
+  dispose(): void {
+    this.keys.dispose();
+    this.values.dispose();
+    this.#free = [];
+  }
+}
+
+/** Cache implementation over a BlockPool. Same values as a plain KVCache
+ *  at every step (parity-gated); only the physical arrangement differs.
+ *  The pool allocates lazily on the first write (KVCache's lazy-alloc
+ *  shape): head count / head dims / dtype come from the first k/v pair,
+ *  so the wiring (maybePageKv) needs no per-model shape plumbing. */
+export class PagedKVCache implements Cache {
+  /** Matches KVCache.STEP: v1's growth granularity is a permutation of
+   *  today's 256-token step into fixed reusable slots, not a new tuning
+   *  axis (--paged-kv-block-size overrides for experiments). */
+  static readonly DEFAULT_BLOCK_SIZE = 256;
+  /** Compiled-decode trace adapters never wrap paged caches (see Cache);
+   *  always unset here. */
+  readonly ropeOffsetArr?: MlxArray;
+  offset = 0;
+  /** Logical block order → physical pool slot. */
+  blockTable: number[] = [];
+  pool: BlockPool | null = null;
+
+  constructor(
+    /** Tokens this cache must be able to hold (prompt + maxTokens). */
+    readonly capacityTokens: number,
+    readonly blockSize: number,
+  ) {}
+
+  get #blockSize(): number {
+    return this.blockSize;
+  }
+
+  /** Blocks needed to hold `n` tokens. */
+  #blocksFor(n: number): number {
+    return Math.ceil(n / this.#blockSize);
+  }
+
+  updateAndFetch(k: MlxArray, v: MlxArray): [MlxArray, MlxArray] {
+    const L = k.shape[2]!;
+    const [, H, , kD] = k.shape as [number, number, number, number];
+    const vD = v.shape[3]!;
+    if (!this.pool)
+      this.pool = new BlockPool({
+        numBlocks: poolBlocksFor(this.capacityTokens, this.blockSize),
+        blockSize: this.blockSize,
+        numKvHeads: H,
+        headDim: kD,
+        vHeadDim: vD,
+        dtype: k.dtype,
+      });
+    // Write, splitting the incoming L along block boundaries. Blocks past
+    // the current tail allocate from the free list as they're reached.
+    let written = 0;
+    while (written < L) {
+      const pos = this.offset + written;
+      const bi = Math.floor(pos / this.#blockSize);
+      const within = pos % this.#blockSize;
+      while (this.blockTable.length <= bi) this.blockTable.push(this.pool.alloc());
+      const l = Math.min(this.#blockSize - within, L - written);
+      // Full-range pieces skip the slice (fresh-view slice would be a
+      // gratuitous op); partial pieces slice the [written, written+l) rows.
+      if (l === L && written === 0) {
+        this.pool.writeBlock(this.blockTable[bi]!, k, v, within);
+      } else {
+        const kp = k.slice([0, 0, written, 0], [1, H, written + l, kD]);
+        const vp = v.slice([0, 0, written, 0], [1, H, written + l, vD]);
+        this.pool.writeBlock(this.blockTable[bi]!, kp, vp, within);
+        kp.dispose();
+        vp.dispose();
+      }
+      written += l;
+    }
+    this.offset += L;
+
+    // Fetch: gather occupied blocks contiguous, slice to the live prefix.
+    const nb = this.#blocksFor(this.offset);
+    const [gk, gv] = this.pool.gather(this.blockTable.slice(0, nb));
+    const keys = gk.slice([0, 0, 0, 0], [1, H, this.offset, kD]);
+    const values = gv.slice([0, 0, 0, 0], [1, H, this.offset, vD]);
+    gk.dispose();
+    gv.dispose();
+    return [keys, values];
+  }
+
+  /** Same mask policy as KVCache.makeMask (it reads only offset). */
+  makeMask(N: number, windowSize: number | null): Mask {
+    if (N === 1) return { mode: "", arr: null };
+    if (windowSize === null) return { mode: "causal", arr: null };
+    if (this.offset === 0 && N <= windowSize) return { mode: "causal", arr: null };
+    return { mode: "array", arr: createCausalMask(N, this.offset, windowSize) };
+  }
+
+  /** Pool tensors — what a prefill-chunk boundary must materialize.
+   *  NOT prompt-cache-compatible in v1 (cloneKvCaches never sees paged
+   *  caches; the serve lane bypasses take/put for paged requests). */
+  state(): MlxArray[] {
+    return this.pool ? [this.pool.keys, this.pool.values] : [];
+  }
+
+  isTrimmable(): boolean {
+    return true;
+  }
+
+  /** Rewind the last n tokens and return now-unoccupied tail blocks to
+   *  the free list. Stale bytes past offset are never read (fetch slices
+   *  to offset; writes cover exact ranges) — the KVCache padding
+   *  invariant, block-shaped. */
+  trim(n: number): void {
+    this.offset = Math.max(0, this.offset - n);
+    const need = this.#blocksFor(this.offset);
+    while (this.blockTable.length > need) this.pool!.free(this.blockTable.pop()!);
+  }
+
+  dispose(): void {
+    // The cache owns its pool in v1 (one sequence per pool).
+    this.pool?.dispose();
+    this.pool = null;
+    this.blockTable = [];
+    this.offset = 0;
+  }
+}
+
+/** Size a pool for one request: capacity tokens rounded up to whole
+ *  blocks. Exported for the wiring in generate.ts (maybePageKv). */
+export function poolBlocksFor(capacityTokens: number, blockSize: number): number {
+  return Math.max(1, Math.ceil(capacityTokens / blockSize));
+}

package/src/model/qwen3-delta.ts

@@ -22,6 +22,7 @@
 import { MlxArray } from "../mlx/array";
 import { Dtype } from "../mlx/ffi";
 import { MetalKernel } from "../mlx/metal-kernel";
+import { CompiledFunction } from "../mlx/compile";
 import * as ops from "../mlx/ops";
 import type { Cache, Mask } from "./gemma4-base";
 
@@ -118,21 +119,33 @@ function getKernel(): MetalKernel {
 /** compute_g: exp(-exp(A_log_f32) * softplus(a + dt_bias)). Output is f32
  *  (the f32 exp(A_log) promotes the bf16 softplus term) — the kernel reads g
  *  as float, so this dtype is load-bearing. */
+let _computeGFn: CompiledFunction | null = null;
 export function computeG(aLog: MlxArray, a: MlxArray, dtBias: MlxArray): MlxArray {
-  const aLogF = aLog.astype(Dtype.float32);
-  const expA = ops.exp(aLogF); // f32 [Hv]
-  aLogF.dispose();
-  const negExpA = ops.neg(expA); // f32 [Hv]
-  expA.dispose();
-  const adt = ops.add(a, dtBias); // bf16 [B,S,Hv]
-  const sp = ops.softplus(adt); // bf16 [B,S,Hv]
-  adt.dispose();
-  const prod = ops.mul(negExpA, sp); // f32 [B,S,Hv]
-  negExpA.dispose();
-  sp.dispose();
-  const g = ops.exp(prod); // f32 [B,S,Hv]
-  prod.dispose();
-  return g;
+  // mlx-lm gated_delta.py: `@partial(mx.compile, shapeless=True) compute_g` — the
+  // WHOLE chain (exp/neg/softplus/mul/exp) is ONE fused kernel there, not five
+  // separate dispatches. Traced once, replayed.
+  if (!_computeGFn) {
+    _computeGFn = new CompiledFunction((inp) => {
+      const aL = inp[0]!, aa = inp[1]!, dt = inp[2]!;
+      const aLogF = aL.astype(Dtype.float32);
+      const expA = ops.exp(aLogF); aLogF.dispose();
+      const negExpA = ops.neg(expA); expA.dispose();
+      const adt = ops.add(aa, dt);
+      // mlx nn.softplus(x) = mx.logaddexp(x, 0), where 0 is a weak Python int
+      // (no cast node). ops.softplus builds the 0 via scalarLike (fromFloat32 →
+      // f32→bf16 AsType). Build the 0 at adt's dtype from raw bytes so there is
+      // no cast — matching mlx's graph exactly.
+      const zbytes = adt.dtype === Dtype.float32
+        ? new Uint8Array([0, 0, 0, 0])
+        : new Uint8Array([0, 0]);
+      const zero = MlxArray.fromBytesCopy(zbytes, [], adt.dtype);
+      const sp = ops.logaddexp(adt, zero); adt.dispose(); zero.dispose();
+      const prod = ops.mul(negExpA, sp); negExpA.dispose(); sp.dispose();
+      const g = ops.exp(prod); prod.dispose();
+      return [g];
+    });
+  }
+  return _computeGFn.apply([aLog, a, dtBias])[0]!;
 }
 
 /** gated_delta_update (use_kernel path): returns [y, newState].
@@ -185,6 +198,19 @@ export class SSMCache implements Cache {
   conv: MlxArray | null = null;
   recurrent: MlxArray | null = null;
   offset = 0;
+  /** 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
+   *  server's bookkeeping; ours lives here because prompt-cache put()
+   *  keys on EXACT coverage (#putOrDispose offset === tokens.length) and
+   *  merged rows carry different totals. Seeded by mergeRows, advanced in
+   *  lockstep by advance(), row-filtered by filter(). */
+  offsets: number[] | null = null;
+
+  /** Row `i`'s own token coverage (per-row when batched, `offset` serial). */
+  rowOffset(i: number): number {
+    return this.offsets ? this.offsets[i]! : this.offset;
+  }
 
   /** Linear-attn layers don't go through the KV update path. */
   updateAndFetch(): [MlxArray, MlxArray] {
@@ -201,6 +227,10 @@ export class SSMCache implements Cache {
 
   advance(n: number): void {
     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
+    // seeds (each row's prompt length).
+    if (this.offsets) for (let i = 0; i < this.offsets.length; i++) this.offsets[i]! += n;
   }
 
   state(): MlxArray[] {
@@ -230,7 +260,8 @@ export class SSMCache implements Cache {
   // left-padding, so the batched twins of mergeKVRows/filterKVRows collapse
   // to B-axis concat/take. `offset` is bookkeeping only for SSM layers (RoPE
   // reads the FULL-attention caches; the linear-attn forward never reads it),
-  // so merged rows carrying different token counts share the max.
+  // so merged rows carrying different token counts share the max; the exact
+  // per-row counts live in `offsets` (extraction's coverage key).
 
   /** Merge a fully-prefilled solo row into a running batched cache (`prev`
    *  null ⇒ the solo row becomes the batch). The returned cache owns fresh
@@ -253,6 +284,9 @@ export class SSMCache implements Cache {
       solo.recurrent = null;
     }
     out.offset = Math.max(prev?.offset ?? 0, solo.offset);
+    // Seed per-row coverage: an adopted serial `prev` (offsets null) is one
+    // row whose exact count is its scalar offset.
+    out.offsets = [...(prev ? prev.offsets ?? [prev.offset] : []), solo.offset];
     return out;
   }
 
@@ -267,5 +301,32 @@ export class SSMCache implements Cache {
     this.recurrent.dispose();
     this.conv = conv;
     this.recurrent = recurrent;
+    if (this.offsets) this.offsets = keep.map((i) => this.offsets![i]!);
+  }
+
+  /** 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 =
+   *  the row's OWN coverage, not the shared max. */
+  extractRow(i: number): SSMCache {
+    if (!this.conv || !this.recurrent)
+      throw new Error("SSMCache.extractRow: empty cache");
+    const cut = (a: MlxArray): MlxArray => {
+      const lo = a.shape.map(() => 0);
+      lo[0] = i;
+      const hi = [...a.shape];
+      hi[0] = i + 1;
+      const view = a.slice(lo, hi);
+      const own = ops.contiguous(view);
+      view.dispose();
+      return own;
+    };
+    const out = new SSMCache();
+    out.conv = cut(this.conv);
+    out.recurrent = cut(this.recurrent);
+    out.offset = this.rowOffset(i);
+    return out;
   }
 }