mlx-bun 0.0.10 → 0.0.11

package/src/spec/dspark/generate-dflash.ts

@@ -15,12 +15,22 @@ import * as ops from "../../mlx/ops";
 import { Gemma4Model } from "../../model/gemma4";
 import type { DflashDrafter } from "./module-dflash";
 import { KeyStream, processLogits, probsOf, probAtToken, sampleToken, sampleResidual, type DSparkSampleConfig } from "./sample";
+import type { ConfSample } from "./calibration";
 
 export interface DflashGenOptions {
   gamma?: number;
   maxTokens?: number;
   eosTokenIds?: number[];
   sample?: DSparkSampleConfig;
+  /** Confidence-scheduled draft-length pruning (Alg 1). Defaults to the
+   *  checkpoint's STS thresholds (drafter.cfg.sts) when calibrated; pass
+   *  minConf for a uniform manual override; pass thresholds:[] to disable. */
+  thresholds?: number[];
+  minConf?: number;
+  /** Calibration hook (scripts/dspark-calibrate.ts): called once per draft
+   *  round with that round's per-position (confidence, accepted) samples,
+   *  AFTER kAccept is known. Not used by ordinary decode. */
+  onRound?: (samples: ConfSample[]) => void;
 }
 export interface DflashStats {
   emitted: number; drafted: number; accepted: number; targetCalls: number;
@@ -32,19 +42,26 @@ export interface DflashResult { tokens: number[]; stats: DflashStats }
  *  hidden [1,L,H] AND the m-layer context [1,L,m*H] (tapLayers concatenated on
  *  the feature axis). */
 function forwardTapped(model: Gemma4Model, ids: MlxArray, caches: import("../../model/gemma4").Cache[], tapLayers: number[]): { finalH: MlxArray; ctxML: MlxArray } {
-  model.hiddenTap = { layers: new Set(tapLayers), captured: new Map() };
-  let finalH: MlxArray;
+  const cap = new Map<number, MlxArray>();
+  model.hiddenTap = { layers: new Set(tapLayers), captured: cap };
+  let finalH: MlxArray | null = null;
   let ctxML: MlxArray;
   try {
     finalH = model.forwardHidden(ids, caches);
-    const cap = model.hiddenTap.captured;
     const perLayer = tapLayers.map((li) => { const a = cap.get(li); if (!a) throw new Error(`layer ${li} not captured`); return a; });
     ctxML = ops.concatAxis(perLayer, 2); // [1,L,m*H]
     for (const [, a] of cap) a.dispose();
+    cap.clear(); // consumed — the finally must not double-dispose
+    const out = { finalH, ctxML };
+    finalH = null; // ownership returned to the caller
+    return out;
   } finally {
+    // Free partial captures + an orphaned finalH on the throw path (on success
+    // both are already gone).
+    finalH?.dispose();
+    for (const [, a] of cap) a.dispose();
     model.hiddenTap = null;
   }
-  return { finalH, ctxML };
 }
 
 function sampleFromHidden(model: Gemma4Model, hidden: MlxArray, pos: number, cfg: DSparkSampleConfig | null, keys: KeyStream): number {
@@ -90,14 +107,22 @@ function verifySampling(model: Gemma4Model, vHidden: MlxArray, draftLogits: MlxA
   let verdict: Verdict | null = null;
   for (let k = 0; k < gamma; k++) {
     const x = drafts[k]!;
-    const ps = processLogits(slice(tl, k), cfg); const P = probsOf(ps); ps.dispose();
-    const qs = processLogits(slice(draftLogits, k), cfg); const Q = probsOf(qs); qs.dispose();
+    // processLogits does not consume its input — bind each sliced row and
+    // dispose it after (an inline slice(...) here leaked one [1,V] per position).
+    const tRow = slice(tl, k);
+    const ps = processLogits(tRow, cfg); tRow.dispose(); const P = probsOf(ps); ps.dispose();
+    const qRow = slice(draftLogits, k);
+    const qs = processLogits(qRow, cfg); qRow.dispose(); const Q = probsOf(qs); qs.dispose();
     const pX = probAtToken(P, x), qX = probAtToken(Q, x);
     if (u[k]! < Math.min(1, qX > 0 ? pX / qX : 1)) { P.dispose(); Q.dispose(); continue; }
     const corr = sampleResidual(P, Q, keys.next()); P.dispose(); Q.dispose();
     verdict = { kAccept: k, emit: corr }; break;
   }
-  if (!verdict) { const bs = processLogits(slice(tl, gamma), cfg); verdict = { kAccept: gamma, emit: sampleToken(bs, keys.next()) }; bs.dispose(); }
+  if (!verdict) {
+    const bRow = slice(tl, gamma);
+    const bs = processLogits(bRow, cfg); bRow.dispose();
+    verdict = { kAccept: gamma, emit: sampleToken(bs, keys.next()) }; bs.dispose();
+  }
   tl.dispose();
   return verdict;
 }
@@ -114,6 +139,12 @@ export function dflashGenerate(model: Gemma4Model, drafter: DflashDrafter, promp
   const out: number[] = [];
   let confSum = 0, rounds = 0, acceptLenSum = 0;
   let H_ctx: MlxArray | null = null;
+  // Round-local tensors, hoisted so the finally frees them on a mid-round
+  // throw (verify, sampling, concat); each is nulled at its normal disposal
+  // — the serve-loop discipline (src/spec/serve-loop.ts).
+  let vHidden: MlxArray | null = null;
+  let vCtxML: MlxArray | null = null;
+  let blockLogits: MlxArray | null = null;
 
   try {
     // 1. prefill (tapped): seed H_ctx with the prompt's m-layer hiddens.
@@ -132,31 +163,49 @@ export function dflashGenerate(model: Gemma4Model, drafter: DflashDrafter, promp
     out.push(next); stats.emitted++;
 
     const tDecode = performance.now();
+    // Alg 1 thresholds: explicit option > checkpoint STS calibration > none.
+    const thresholds = options.thresholds ?? drafter.cfg.sts?.thresholds;
     outer: while (stats.emitted < maxTokens) {
-      // 2a. draft against the current (growing) H_ctx.
-      const block = drafter.forwardInfer(model, H_ctx!, next, gamma, { sample: sampleCfg ?? undefined, keys });
+      // 2a. draft against the current (growing) H_ctx. The scheduler may prune
+      // the block short (d ≤ γ) — verify then covers exactly d positions.
+      const block = drafter.forwardInfer(model, H_ctx!, next, gamma, {
+        sample: sampleCfg ?? undefined, keys, thresholds, minConf: options.minConf,
+      });
       const drafts = block.tokens;
-      stats.drafted += gamma;
+      blockLogits = block.draftLogits ?? null;
+      const d = drafts.length;
+      stats.drafted += d;
       for (const c of block.conf) confSum += c;
 
       // 2b. verify (tapped) — get accepted tokens' m-layer hiddens to append.
-      const vIds = ops.fromInt32([next, ...drafts], [1, gamma + 1]);
-      const { finalH: vHidden, ctxML: vCtxML } = forwardTapped(model, vIds, caches, tapLayers);
+      const vIds = ops.fromInt32([next, ...drafts], [1, d + 1]);
+      const ft = forwardTapped(model, vIds, caches, tapLayers);
+      vHidden = ft.finalH;
+      vCtxML = ft.ctxML;
       vIds.dispose();
       stats.targetCalls++;
 
       const { kAccept, emit } = sampleCfg
-        ? verifySampling(model, vHidden, block.draftLogits, drafts, gamma, sampleCfg, keys)
-        : verifyGreedy(model, vHidden, drafts, gamma);
-      block.draftLogits.dispose();
+        ? verifySampling(model, vHidden, blockLogits!, drafts, d, sampleCfg, keys)
+        : verifyGreedy(model, vHidden, drafts, d);
+      blockLogits?.dispose();
+      blockLogits = null;
       stats.accepted += kAccept; rounds++; acceptLenSum += kAccept + 1;
 
+      // Calibration hook (scripts/dspark-calibrate.ts): report this round's
+      // per-position (confidence, accepted) outcomes now that kAccept is known.
+      if (options.onRound) {
+        const roundSamples: ConfSample[] = [];
+        for (let k = 0; k < drafts.length; k++) roundSamples.push({ pos: k, conf: block.conf[k]!, accepted: k < kAccept });
+        options.onRound(roundSamples);
+      }
+
       // 2c. emit accepted drafts
       for (let k = 0; k < kAccept; k++) {
         const isEos = eosTokenIds.includes(drafts[k]!);
         if (!isEos) out.push(drafts[k]!);
         stats.emitted++;
-        if (isEos || stats.emitted >= maxTokens) { vHidden.dispose(); vCtxML.dispose(); break outer; }
+        if (isEos || stats.emitted >= maxTokens) break outer; // finally frees vHidden/vCtxML
       }
       const emitIsEos = eosTokenIds.includes(emit);
       if (!emitIsEos) out.push(emit);
@@ -164,21 +213,24 @@ export function dflashGenerate(model: Gemma4Model, drafter: DflashDrafter, promp
 
       // 2d. append accepted stream's m-layer hiddens (anchor + kAccept drafts =
       // vCtxML[:, 0..kAccept]) to H_ctx; never the rejected tips.
-      const add = vCtxML.slice([0, 0, 0], [1, kAccept + 1, vCtxML.shape[2]!]);
+      const add = vCtxML!.slice([0, 0, 0], [1, kAccept + 1, vCtxML!.shape[2]!]);
       const grown = ops.concatAxis([H_ctx!, add], 1);
-      H_ctx!.dispose(); add.dispose(); vCtxML.dispose();
+      H_ctx!.dispose(); add.dispose(); vCtxML!.dispose(); vCtxML = null;
       H_ctx = grown;
 
       // 2e. roll back rejected KV tips (bypass trim, past-window safe)
-      if (kAccept < gamma) { const n = gamma - kAccept; for (const c of caches) c.trim(n, true); }
+      if (kAccept < d) { const nRej = d - kAccept; for (const c of caches) c.trim(nRej, true); }
 
-      vHidden.dispose();
+      vHidden!.dispose(); vHidden = null;
       if (emitIsEos || stats.emitted >= maxTokens) break;
       next = emit;
     }
     stats.decodeMs = performance.now() - tDecode;
   } finally {
     H_ctx?.dispose();
+    vHidden?.dispose();
+    vCtxML?.dispose();
+    blockLogits?.dispose();
     for (const c of caches) c.dispose();
   }
   return finalize(stats, out, rounds, acceptLenSum, confSum);

package/src/spec/dspark/loader.ts

@@ -0,0 +1,23 @@
+// Central DSpark checkpoint loader — the variant dispatch the modules never
+// had (scripts imported their class directly; the serve path needs one entry
+// point). Reads dspark.json and returns the matching drafter:
+//   variant "dspark"           — the faithful KV-injection module (canonical)
+//   variant "dflash" (legacy)  — same module, pre-rename stamp
+//   variant absent             — a v1 single-vector checkpoint (module.ts,
+//                                superseded research baseline) — refused for
+//                                serving with a pointer at the v2 trainer.
+
+import { readFileSync, existsSync } from "node:fs";
+import { join } from "node:path";
+import { DflashDrafter } from "./module-dflash";
+
+export function loadDsparkDrafter(dir: string): DflashDrafter {
+  const metaPath = join(dir, "dspark.json");
+  if (!existsSync(metaPath)) throw new Error(`no dspark.json in ${dir} — not a DSpark checkpoint`);
+  const meta = JSON.parse(readFileSync(metaPath, "utf8")) as { variant?: string };
+  if (meta.variant === "dspark" || meta.variant === "dflash") return DflashDrafter.load(dir);
+  throw new Error(
+    `${dir} is a v1 single-vector DSpark checkpoint (variant=${meta.variant ?? "none"}) — ` +
+      `superseded, not serveable; retrain with scripts/dspark-train-dflash.ts`,
+  );
+}

package/src/spec/dspark/module-dflash.ts

@@ -23,6 +23,21 @@ import { processLogits, sampleToken, KeyStream, type DSparkSampleConfig } from "
 import { mkdirSync, writeFileSync, readFileSync, existsSync } from "node:fs";
 import { join } from "node:path";
 
+/** STS calibration (§3.2.1) — per-position confidence thresholds fit on real
+ *  (confidence, accepted) verify outcomes (src/spec/dspark/calibration.ts,
+ *  scripts/dspark-calibrate.ts). Consumed by the confidence-scheduled
+ *  draft-length pruning (Alg 1, single-user form) in forwardInfer. Stored
+ *  inside DflashConfig so it round-trips through dspark.json for free.
+ *  ABSENT ⟹ no pruning — an uncalibrated checkpoint drafts exactly as before. */
+export interface StsCalibration {
+  /** thresholds[k] = min confidence to KEEP draft position k. length γ. */
+  thresholds: number[];
+  /** Target acceptance precision the thresholds were fit for (e.g. 0.5). */
+  target: number;
+  /** Verify outcomes the fit saw (provenance). */
+  samples?: number;
+}
+
 export interface DflashConfig {
   gamma: number;
   dDraft: number;
@@ -34,6 +49,10 @@ export interface DflashConfig {
   /** Target layer indices tapped for H_ctx (Eq 2). nLayers-of-target..sentinel.
    *  For e4b (42 layers): index 0..41 = post-layer residual, 42 = post-finalNorm. */
   tapLayers: number[];
+  /** Sequential head: "markov" (Eq 5, default) or "rnn" (Eq 6). */
+  seqHead?: "markov" | "rnn";
+  /** STS calibration (§3.2.1); absent = no confidence pruning. */
+  sts?: StsCalibration;
 }
 
 export const DEFAULT_DFLASH_CONFIG: DflashConfig = {
@@ -49,8 +68,30 @@ export const DEFAULT_DFLASH_CONFIG: DflashConfig = {
 export interface TargetDims { hiddenSize: number; vocabSize: number; eps: number }
 
 export interface DflashTrainOut { draftLogits: MlxArray; conf: MlxArray }
-export interface DflashDraftBlock { tokens: number[]; conf: number[]; draftLogits: MlxArray }
-export interface DflashDraftOpts { sample?: DSparkSampleConfig; keys?: KeyStream }
+/** tokens/conf may be SHORTER than the requested γ when confidence pruning
+ *  fires (Alg 1); always ≥1. draftLogits covers exactly tokens.length
+ *  positions — ABSENT when the caller passed collectLogits:false (the loop
+ *  never materialized/concatenated the per-position logits). */
+export interface DflashDraftBlock { tokens: number[]; conf: number[]; draftLogits?: MlxArray }
+export interface DflashDraftOpts {
+  sample?: DSparkSampleConfig;
+  keys?: KeyStream;
+  /** Confidence-scheduled draft-length pruning (Alg 1, single-user form):
+   *  position k is DROPPED (and the block truncated there) when its predicted
+   *  acceptance c_k < thresholds[k] (per-position, STS-calibrated) or < minConf
+   *  (uniform manual override). Position 0 is always kept (a source must
+   *  return ≥1 token). Both absent ⟹ fixed-γ drafting (pre-scheduler
+   *  behavior). Losslessness is invariant — pruning only changes how many
+   *  positions the target VERIFIES, never what is emitted. */
+  thresholds?: number[];
+  minConf?: number;
+  /** Materialize+concat the per-position draft logits into the returned
+   *  DflashDraftBlock.draftLogits. Default true (the standalone
+   *  dflashGenerate sampling-verify path reads it). The serve-loop draft()
+   *  path runs its own verify lm-head and never reads it — pass false there
+   *  to skip the dead concat + a host-unrelated GPU alloc every round. */
+  collectLogits?: boolean;
+}
 
 const CDT = Dtype.float32;
 const NEG = -1e9;
@@ -126,6 +167,15 @@ export class DflashDrafter {
     set("markov.w2", zerosArray([r, V]));   // W2=0 → starts as pure parallel DFlash
     set("conf.w", zerosArray([dDraft + r, 1]));
     set("conf.b", zerosArray([1]));
+    // rnn.* MUST be initialized last (after every markov/shared param) so the
+    // seeded rng stream hands out IDENTICAL draws to a markov-only init up to
+    // this point — this is what makes the init-equivalence gate hold (an rnn
+    // drafter and a markov drafter, same seed, are bit-identical at t=0).
+    if (cfg.seqHead === "rnn") {
+      set("rnn.wH", normalArray(rng, [r, r], 0.02));
+      set("rnn.bH", zerosArray([r]));
+      set("rnn.wO", zerosArray([r, V])); // zero-init → starts as pure DFlash, mirrors markov.w2
+    }
     return d;
   }
 
@@ -261,6 +311,61 @@ export class DflashDrafter {
     const B = ops.matmul(emb1, this.get("markov.w2")); // [A,γ,V]
     return { B, emb1 };
   }
+
+  /** RNN sequential head — the ELMAN variant. VERIFIED-DIVERGENT from paper
+   *  Eq 6 (2026-07-06 audit vs arXiv:2607.05147): the paper's cell is GATED
+   *  over z_k=[s_{k−1}; W₁[x_{k−1}]; h_k] (backbone hidden included!) with a
+   *  fused (2r+d)→3r projection, s_k=σ(W_g z)⊙s_{k−1}+(1−σ)⊙tanh(W_c z),
+   *  B_k=W₂ᵀtanh(W_o z) — DeepSpec's RNNHead is that form. This ungated
+   *  tanh recurrence over the shared token embedding (no h_k) is kept as a
+   *  documented VARIANT ([[dont-delete-optionality-on-one-measurement]]):
+   *  the paper reports the RNN head's gains over Markov as "marginal" and
+   *  production uses Markov (our default). Add the gated cell as
+   *  seqHead:"gru" if it ever earns a run; details in the design doc's
+   *  "DeepSpec ground truth" section. */
+  #rnn(prevToks: MlxArray): { B: MlxArray; emb1: MlxArray } {
+    const { gamma: G, markovRank: r } = this.cfg;
+    const A = prevToks.shape[0]!;
+    const emb1 = ops.takeAxis(this.get("markov.w1"), prevToks, 0); // [A,γ,r] (shared embedding, no 2nd table)
+    const wH = this.get("rnn.wH"), bH = this.get("rnn.bH"), wO = this.get("rnn.wO");
+    let s: MlxArray | null = zerosArray([A, r]); // s_{-1} = 0
+    const bSlices: MlxArray[] = [];
+    const pending: MlxArray[] = []; // iteration temps, for the catch (forwardInfer pattern)
+    try {
+      for (let k = 0; k < G; k++) {
+        const e1S = emb1.slice([0, k, 0], [A, k + 1, r]); pending.push(e1S);
+        const e1k = ops.reshape(e1S, [A, r]); pending.push(e1k);
+        e1S.dispose();
+        const sh = ops.matmul(s!, wH); pending.push(sh);
+        s!.dispose(); s = null;
+        let sNext = ops.add(sh, e1k); pending.push(sNext);
+        sh.dispose(); e1k.dispose();
+        sNext = disposing(sNext, ops.add(sNext, bH)); pending.push(sNext);
+        sNext = disposing(sNext, ops.tanh(sNext)); pending.push(sNext);
+        s = sNext; // s_k
+        const bk = ops.matmul(s, wO); pending.push(bk); // [A,V]
+        bSlices.push(ops.reshape(bk, [A, 1, this.dims.vocabSize])); bk.dispose();
+        pending.length = 0;
+      }
+      s!.dispose(); s = null;
+      const B = ops.concatAxis(bSlices, 1); // [A,γ,V]
+      for (const b of bSlices) b.dispose();
+      return { B, emb1 };
+    } catch (err) {
+      // Mid-recurrence throw: free the carried state, the collected slices,
+      // the current iteration's temps, and the emb1 the caller never received
+      // (dispose() idempotence makes over-disposal inert).
+      for (const a of pending) a.dispose();
+      for (const b of bSlices) b.dispose();
+      s?.dispose();
+      emb1.dispose();
+      throw err;
+    }
+  }
+
+  #seqHead(prevToks: MlxArray): { B: MlxArray; emb1: MlxArray } {
+    return this.cfg.seqHead === "rnn" ? this.#rnn(prevToks) : this.#markov(prevToks);
+  }
   #confidence(block: MlxArray, emb1: MlxArray): MlxArray {
     const inp = ops.concatAxis([block, emb1], 2);
     let z = ops.matmul(inp, this.get("conf.w")); inp.dispose();
@@ -275,23 +380,49 @@ export class DflashDrafter {
    *  hCtx [A,Lctx,m*H], ctxMask [A,Lctx]|null, anchorEmb [A,H], prevToks [A,γ]. */
   forwardTrain(model: Gemma4Model, hCtx: MlxArray, ctxMask: MlxArray | null, anchorEmb: MlxArray, prevToks: MlxArray): DflashTrainOut {
     const A = hCtx.shape[0]!;
-    const block = this.#backbone(hCtx, anchorEmb, ctxMask, A);
-    const U = this.#baseLogits(model, block);
-    const Uf = U.astype(CDT); U.dispose();
-    const { B, emb1 } = this.#markov(prevToks);
-    const draftLogits = disposing(Uf, ops.add(Uf, B)); B.dispose();
-    const conf = this.#confidence(block, emb1);
-    block.dispose(); emb1.dispose();
-    return { draftLogits, conf };
+    let block: MlxArray | null = null, Uf: MlxArray | null = null;
+    let B: MlxArray | null = null, emb1: MlxArray | null = null;
+    try {
+      block = this.#backbone(hCtx, anchorEmb, ctxMask, A);
+      const U = this.#baseLogits(model, block);
+      Uf = U.astype(CDT); U.dispose();
+      ({ B, emb1 } = this.#seqHead(prevToks));
+      const draftLogits = ops.add(Uf, B);
+      Uf.dispose(); Uf = null; B.dispose(); B = null;
+      const conf = this.#confidence(block, emb1);
+      block.dispose(); block = null; emb1.dispose(); emb1 = null;
+      return { draftLogits, conf };
+    } catch (err) {
+      // Free whatever a mid-forward throw left live (each slot is nulled the
+      // moment its normal disposal ran — no double-free).
+      block?.dispose(); Uf?.dispose(); B?.dispose(); emb1?.dispose();
+      throw err;
+    }
   }
 
   /** Inference: parallel backbone once, then sequential Markov sampling.
-   *  hCtx [1,Lctx,m*H] (full current context), anchorTok the bonus token. */
+   *  hCtx [1,Lctx,m*H] (full current context), anchorTok the bonus token.
+   *
+   *  Host-sync discipline (docs/investigations/dspark-handoff.md item 3): the
+   *  greedy path's token recurrence stays ON-DEVICE across positions — the
+   *  argmax result [1] uint32 feeds directly into the next position's
+   *  takeAxis, no itemUint32 in the loop. Per-position token arrays are
+   *  collected and resolved to host numbers with exactly ONE concat +
+   *  toIntTokens() after the loop. Confidence is the one read the pruning
+   *  ALGORITHM inherently needs per position (Alg 1 must decide whether to
+   *  stop NOW) — so it stays per-position whenever thresholds/minConf are
+   *  actually active; when pruning is inactive there is no such need, so
+   *  confidence is deferred the same way (collect [1] arrays, one concat +
+   *  one toFloat32 after the loop). The sampling path already forces a host
+   *  read every position (sampleToken draws from a device RNG key) — that
+   *  path is unchanged. */
   forwardInfer(model: Gemma4Model, hCtx: MlxArray, anchorTok: number, gamma: number, opts: DflashDraftOpts = {}): DflashDraftBlock {
     const { dDraft, markovRank: r } = this.cfg;
     const V = this.dims.vocabSize;
     const sample = opts.sample && opts.sample.temperature > 0 ? opts.sample : null;
     const keys = opts.keys ?? new KeyStream(opts.sample?.seed ?? 0);
+    const collectLogits = opts.collectLogits ?? true;
+    const pruningActive = opts.thresholds !== undefined || opts.minConf !== undefined;
 
     const ids = ops.fromInt32([anchorTok], [1, 1]);
     const ae3 = model.embed.encode(ids); ids.dispose();
@@ -302,30 +433,134 @@ export class DflashDrafter {
     const Uf = U.astype(CDT); U.dispose();
 
     const w1 = this.get("markov.w1"), w2 = this.get("markov.w2");
-    const tokens: number[] = [], conf: number[] = [];
-    const perPos: MlxArray[] = [];
-    let prevTok = anchorTok;
-    for (let k = 0; k < gamma; k++) {
-      const Uk = ops.reshape(Uf.slice([0, k, 0], [1, k + 1, V]), [1, V]);
-      const prevIdx = ops.fromInt32([prevTok], [1]);
-      const e1 = ops.takeAxis(w1, prevIdx, 0); prevIdx.dispose();
-      const Bk = ops.matmul(e1, w2);
-      const logitsK = ops.add(Uk, Bk); Uk.dispose(); Bk.dispose();
-      let tok: number;
-      if (sample) { const sc = processLogits(logitsK, sample); tok = sampleToken(sc, keys.next()); sc.dispose(); }
-      else { const am = ops.argmaxAxis(logitsK, -1); tok = ops.itemUint32(am); am.dispose(); }
-      perPos.push(ops.reshape(logitsK, [1, 1, V])); logitsK.dispose();
-      const hk = ops.reshape(block.slice([0, k, 0], [1, k + 1, dDraft]), [1, dDraft]);
-      const ci = ops.concatAxis([hk, e1], 1); hk.dispose(); e1.dispose();
-      let cz = ops.matmul(ci, this.get("conf.w")); ci.dispose();
-      cz = disposing(cz, ops.add(cz, this.get("conf.b")));
-      cz = disposing(cz, ops.sigmoid(cz));
-      conf.push(cz.toFloat32()[0]!); cz.dispose();
-      tokens.push(tok); prevTok = tok;
+    const isRnn = this.cfg.seqHead === "rnn";
+    const wH = isRnn ? this.get("rnn.wH") : null;
+    const bH = isRnn ? this.get("rnn.bH") : null;
+    const wO = isRnn ? this.get("rnn.wO") : null;
+    let s: MlxArray | null = isRnn ? zerosArray([1, r]) : null; // s_{-1} = 0, carried across k
+    const conf: number[] = [];
+    const perPos: MlxArray[] = []; // per-position [1,1,V] draft logits (collectLogits only)
+    const tokArrs: MlxArray[] = []; // per-position [1] device token ids, greedy path only
+    const confArrs: MlxArray[] = []; // per-position [1] device confidences, deferred (no pruning) only
+    let tokens: number[] = [];
+    let prevIdx: MlxArray = ops.fromInt32([anchorTok], [1]); // device token id fed to takeAxis
+    let ownPrevIdx = true; // false once prevIdx aliases a tokArrs/argmax entry we must not double-dispose
+    // Iteration-local temps register here at creation so the catch can reach
+    // them on a mid-flight throw; drained (length=0) at the bottom of each
+    // iteration once everything in it was disposed on the normal path.
+    // dispose() idempotence makes catch-side re-disposal of the already-freed
+    // ones inert.
+    const pending: MlxArray[] = [];
+    try {
+      for (let k = 0; k < gamma; k++) {
+        const UkS = Uf.slice([0, k, 0], [1, k + 1, V]); pending.push(UkS);
+        const Uk = ops.reshape(UkS, [1, V]); pending.push(Uk);
+        UkS.dispose();
+        const e1 = ops.takeAxis(w1, prevIdx, 0); pending.push(e1); // [1,r], shared embedding either way
+        if (ownPrevIdx) { prevIdx.dispose(); ownPrevIdx = false; } // consumed; prevIdx reassigned below
+        let Bk: MlxArray;
+        if (isRnn) {
+          const sh = ops.matmul(s!, wH!); pending.push(sh);
+          s!.dispose(); s = null;
+          let sNext = ops.add(sh, e1); pending.push(sNext);
+          sh.dispose();
+          sNext = disposing(sNext, ops.add(sNext, bH!)); pending.push(sNext);
+          sNext = disposing(sNext, ops.tanh(sNext)); pending.push(sNext);
+          s = sNext; // s_k, survives to the next iteration (covered by the s slot from here)
+          Bk = ops.matmul(s, wO!); pending.push(Bk);
+        } else {
+          Bk = ops.matmul(e1, w2); pending.push(Bk);
+        }
+        const logitsK = ops.add(Uk, Bk); pending.push(logitsK);
+        Uk.dispose(); Bk.dispose();
+        // prevIdx becomes the SINGLE tracked owner of the next token id the
+        // moment it exists — no separate local for the catch to miss.
+        if (sample) {
+          const sc = processLogits(logitsK, sample); pending.push(sc);
+          const tok = sampleToken(sc, keys.next()); sc.dispose();
+          tokens.push(tok);
+          prevIdx = ops.fromInt32([tok], [1]); ownPrevIdx = true;
+        } else {
+          const am = ops.argmaxAxis(logitsK, -1); // [1] uint32, ON-DEVICE — no itemUint32 here
+          tokArrs.push(am);
+          prevIdx = am; ownPrevIdx = false; // aliases tokArrs[k] — that collection owns it
+        }
+        if (collectLogits) perPos.push(ops.reshape(logitsK, [1, 1, V]));
+        logitsK.dispose();
+        const hkS = block.slice([0, k, 0], [1, k + 1, dDraft]); pending.push(hkS);
+        const hk = ops.reshape(hkS, [1, dDraft]); pending.push(hk);
+        hkS.dispose();
+        const ci = ops.concatAxis([hk, e1], 1); pending.push(ci);
+        hk.dispose(); e1.dispose();
+        let cz = ops.matmul(ci, this.get("conf.w")); pending.push(cz);
+        ci.dispose();
+        cz = disposing(cz, ops.add(cz, this.get("conf.b"))); pending.push(cz);
+        cz = disposing(cz, ops.sigmoid(cz)); pending.push(cz);
+
+        // Alg 1 (single-user): drop position k when its predicted acceptance
+        // is below the calibrated/manual threshold — the target then
+        // verifies a shorter window. Position 0 always survives (sources
+        // return ≥1 token). Pruning inherently needs the confidence value
+        // THIS position, on the host, to decide whether to stop — no way
+        // around that read when active. When inactive, defer it (device
+        // array, resolved once after the loop) since nothing this iteration
+        // depends on its value.
+        const thr = opts.thresholds?.[k] ?? opts.minConf;
+        if (pruningActive) {
+          const cVal = cz.toFloat32()[0]!; cz.dispose();
+          if (thr !== undefined && k > 0 && cVal < thr) {
+            // Drop position k SYMMETRICALLY on both paths so tokens/conf/
+            // draftLogits stay aligned (the DflashDraftBlock contract —
+            // verifySampling slices draftLogits per token). Greedy: pop the
+            // device array (prevIdx aliases it — one dispose via the owning
+            // collection). Sample: pop the host token + free the standalone
+            // fromInt32.
+            if (!sample) tokArrs.pop()!.dispose(); // also frees prevIdx (same array)
+            else { prevIdx.dispose(); tokens.pop(); }
+            ownPrevIdx = false; // dead either way — post-loop/catch must not re-dispose
+            if (collectLogits) perPos.pop()!.dispose(); // logits row for the dropped position
+            break;
+          }
+          conf.push(cVal);
+        } else {
+          confArrs.push(cz);
+        }
+        pending.length = 0; // everything above was disposed (or handed to a collection)
+      }
+      pending.length = 0; // early-break path: the broken iteration's temps are all freed too
+      if (ownPrevIdx) prevIdx.dispose(); // last iteration's fresh sample-path array, unused past the loop
+      s?.dispose(); // final on-device state (loop end or early break — both leave s live)
+      Uf.dispose(); block.dispose();
+
+      if (!sample) {
+        const allToks = ops.concatAxis(tokArrs, 0); // [nKept] uint32, ONE host read below
+        for (const a of tokArrs) a.dispose();
+        tokens = allToks.toIntTokens(); allToks.dispose();
+      }
+      if (!pruningActive && confArrs.length) {
+        const allConf = ops.concatAxis(confArrs, 0); // [nKept,1]
+        for (const a of confArrs) a.dispose();
+        const flat = allConf.toFloat32(); allConf.dispose();
+        for (const v of flat) conf.push(v);
+      }
+      const draftLogits = collectLogits ? ops.concatAxis(perPos, 1) : undefined;
+      if (collectLogits) for (const a of perPos) a.dispose();
+      return { tokens, conf, draftLogits };
+    } catch (err) {
+      // Cleanup on a mid-flight throw — the collections, the carried state,
+      // AND the current iteration's registered temps (`pending`). dispose()
+      // is idempotent (array.ts #disposed guard) so re-disposing an array the
+      // try block already freed (or that aliases one of these collections,
+      // e.g. prevIdx === tokArrs[k]) is inert, never a double-free bug.
+      for (const a of pending) a.dispose();
+      if (ownPrevIdx) prevIdx.dispose();
+      for (const a of tokArrs) a.dispose();
+      for (const a of confArrs) a.dispose();
+      for (const a of perPos) a.dispose();
+      s?.dispose();
+      Uf.dispose(); block.dispose();
+      throw err;
     }
-    Uf.dispose(); block.dispose();
-    const draftLogits = ops.concatAxis(perPos, 1); for (const a of perPos) a.dispose();
-    return { tokens, conf, draftLogits };
   }
 
   save(dir: string): void {
@@ -333,7 +568,9 @@ export class DflashDrafter {
     const tensors: NamedTensor[] = this.names.map((name) => ({ name, array: this.get(name) }));
     writeShardedSafetensors(dir, tensors);
     writeFileSync(join(dir, "dspark.json"), JSON.stringify({
-      kind: "dspark-drafter", variant: "dflash", version: 1,
+      // "dspark" is the canonical variant (this faithful module IS DSpark);
+      // load() accepts the legacy "dflash" stamp from older checkpoints.
+      kind: "dspark-drafter", variant: "dspark", version: 1,
       target_id: this.targetId, config: this.cfg, dims: this.dims, param_names: this.names,
     }, null, 2));
   }
@@ -341,7 +578,8 @@ export class DflashDrafter {
     const meta = JSON.parse(readFileSync(join(dir, "dspark.json"), "utf8")) as {
       variant?: string; config: DflashConfig; dims: TargetDims; target_id: string;
     };
-    if (meta.variant !== "dflash") throw new Error(`${dir} is not a dflash drafter (variant=${meta.variant})`);
+    if (meta.variant !== "dspark" && meta.variant !== "dflash")
+      throw new Error(`${dir} is not a dspark drafter (variant=${meta.variant ?? "none — a v1 single-vector checkpoint?"})`);
     const d = new DflashDrafter(meta.config, meta.dims, meta.target_id);
     const tensors = loadAdapterTensors(join(dir, "model.safetensors"));
     try {
@@ -357,5 +595,9 @@ function buildNames(cfg: DflashConfig): string[] {
   for (let i = 0; i < cfg.nLayers; i++)
     names.push(`bb.${i}.attn_norm`, `bb.${i}.q`, `bb.${i}.k`, `bb.${i}.v`, `bb.${i}.o`, `bb.${i}.ff_norm`, `bb.${i}.gate`, `bb.${i}.up`, `bb.${i}.down`);
   names.push("W_c", "ctx_norm", "tok_proj", "mask_emb", "block_pos", "out_norm", "out_proj", "markov.w1", "markov.w2", "conf.w", "conf.b");
+  // rnn.* MUST be appended LAST — see the matching comment in initFromDims:
+  // this ordering keeps the shared/markov params drawing identical seeded rng
+  // values whether or not the rnn head is present (init-equivalence gate).
+  if (cfg.seqHead === "rnn") names.push("rnn.wH", "rnn.bH", "rnn.wO");
   return names;
 }