mlx-bun 0.2.0 → 0.4.0

package/src/spec/serve-loop.ts

@@ -1,6 +1,6 @@
 // Serve-time speculative decoding — the ONE verify/accept executor behind
-// `serve --draft-model` (docs/design/mlx-lm-tool-parity-plan.md §7; sequenced
-// by docs/design/grammar-spec-batching-integration.md Phase B).
+// `serve --draft-model` (docs/reference/cli.md §7; sequenced
+// by docs/design/batching.md Phase B).
 //
 // Faithful to mlx-lm's speculative_generate_step (generate.py:473-654, read
 // from the oracle venv): per round, draft n tokens from the DraftSource, run
@@ -17,11 +17,9 @@
 //    trimmable mid-generation; a serve endpoint must not 500 mid-stream, so
 //    we STOP SPECULATING and finish the generation with plain single-token
 //    decode (bit-equivalent continuation — the target's own samples).
-//  - Prompt-cache reuse is BYPASSED in v1 (fresh caches per spec request,
-//    cachedTokens=0). mlx-lm composes spec with its LRU prompt cache
-//    (target+draft caches per entry); wiring that through our PromptCache +
-//    SSD tier is a tracked follow-up in the integration plan, not silently
-//    absent.
+//  - Sources with checkpoint support use the injected prefix-cache service.
+//    Target and companion state share the existing RAM/SSD lifecycle. Other
+//    sources begin with fresh state until they supply a checkpoint codec.
 //
 // Grammar × spec (Phase C, the constrained verify walk — novel: NO runtime
 // serves both; mlx-lm has no grammar, oMLX no spec): the drafter runs FREE
@@ -31,7 +29,7 @@
 // so rejected drafts never touch grammar state and no matcher rollback is
 // needed. Grammar termination mid-burst truncates the round (the all--inf
 // guarantee). Gates (no oracle): greedy grammar+spec ≡ greedy grammar-only
-// long-prefix + 100% schema validity (tests/spec-serve.test.ts).
+// long-prefix + 100% schema validity (tests/parity/spec-serve.test.ts).
 //
 // Emission discipline: tokens flow through the caller's onToken ONE AT A
 // TIME, in order (bursts of ≤ n+1 per round) — the server's stop-sequence
@@ -40,70 +38,25 @@
 // false halts the generation mid-burst.
 
 import { MlxArray } from "../mlx/array";
+import { sampleSpeculativeRows } from "../backends/mlx/speculative-round";
 import * as ops from "../mlx/ops";
 import { clearCache } from "../mlx/ffi";
-import { flagOn } from "../flags";
+import { cleanupFailure, disposeResources } from "../engine/resources";
+import type { DisposableResource } from "../contracts/resources";
+import { assertMlxSpeculativeBinding, bindLegacySpeculativeModel, type MlxSpeculativeBinding } from "../backends/mlx/speculative";
 import type { RuntimeModel } from "../model/factory";
+import { cloneKvCaches } from "../kv-store";
+import { disposeAttachments, type CheckpointAttachment } from "../backends/mlx/checkpoint-state";
 import type { Cache } from "../model/gemma4";
-import type { GenerateOptions, GenerateStats } from "../generate";
+import { withModelUsageFlush, withModelWiredLimit, type GenerateOptions, type GenerateStats } from "../generate";
+import { createKvMaintenance } from "../backends/mlx/kv-maintenance";
+import { runtimeConfig, withRuntimeConfig } from "../runtime-config";
 import { makeSampler, makeStepSampler } from "../sampler";
 import type { OnToken } from "../serve/generation-gateway";
 import type { DraftProvider } from "./source";
 
 const PREFILL_CHUNK = 2048;
 
-/** Run one target forward, and when `tapLayers` is set (⟹ a Gemma4 target — a
- *  DSpark-style source), also return the captured multi-layer context
- *  [1,L,m*H] (tapLayers concatenated on the feature axis; index nLayers is the
- *  post-finalNorm sentinel). Non-tapping sources get ctxML=null and never
- *  touch model.hiddenTap. Mirrors generate-dflash.ts forwardTapped. */
-async function forwardMaybeTap(
-  model: RuntimeModel,
-  ids: MlxArray,
-  caches: Cache[],
-  tapLayers: number[] | undefined,
-): Promise<{ hidden: MlxArray; ctxML: MlxArray | null }> {
-  if (!tapLayers) {
-    const asyncModel = model as RuntimeModel & {
-      forwardHiddenAsync?: (
-        ids: MlxArray,
-        caches: Cache[],
-      ) => Promise<MlxArray>;
-    };
-    const hidden = typeof asyncModel.forwardHiddenAsync === "function"
-      ? await asyncModel.forwardHiddenAsync(ids, caches)
-      : model.forwardHidden(ids, caches);
-    return { hidden, ctxML: null };
-  }
-  const m = model as unknown as {
-    hiddenTap: { layers: Set<number>; captured: Map<number, MlxArray> } | null;
-  };
-  const cap = new Map<number, MlxArray>();
-  m.hiddenTap = { layers: new Set(tapLayers), captured: cap };
-  let hidden: MlxArray | null = null;
-  try {
-    hidden = model.forwardHidden(ids, caches);
-    const perLayer = tapLayers.map((li) => {
-      const a = cap.get(li);
-      if (!a) throw new Error(`spec tap: layer ${li} not captured`);
-      return a;
-    });
-    const 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 = { hidden, ctxML };
-    hidden = null; // ownership returned to the caller
-    return out;
-  } finally {
-    // On any throw (forward mid-capture, missing layer, concat), free the
-    // partially-captured tap tensors and the orphaned hidden. On success both
-    // are already gone (cap cleared, hidden nulled).
-    hidden?.dispose();
-    for (const [, a] of cap) a.dispose();
-    m.hiddenTap = null;
-  }
-}
-
 export interface SpecServeExtras {
   drafted: number;
   accepted: number;
@@ -115,12 +68,19 @@ export interface SpecServeExtras {
   forwardsSaved: number;
   /** Per-draft-position counters (index = position within a round's block,
    *  0..γ-1): how many rounds drafted/accepted at that position. Drives the
-   *  Phase-1c per-position acceptance report (scripts/dspark-drafter-ab.ts)
+   *  Phase-1c per-position acceptance report (scripts/dspark.ts ab)
    *  — near-zero cost, always populated. */
   draftedByPos: number[];
   acceptedByPos: number[];
 }
 
+export interface SpecRunServices {
+  readonly prefixCache?: import("../backends/mlx/checkpoint-state").MlxPrefixCache;
+  /** Mounted adapter revision, resolved while the execution lease is held. */
+  readonly cacheNamespace?: string;
+  readonly cloneState?: (caches: Cache[]) => Cache[];
+}
+
 export async function specServeRun(
   model: RuntimeModel,
   provider: DraftProvider,
@@ -128,9 +88,40 @@ export async function specServeRun(
   promptIds: number[],
   options: GenerateOptions & { stopSequences?: string[] },
   onToken: OnToken,
+  services: SpecRunServices = {},
+): Promise<GenerateStats> {
+  return specRun(bindLegacySpeculativeModel(model, provider), numDraftTokens, promptIds, options, onToken, services);
+}
+
+/** The verifier consumes a bound target and source factory, independent of
+ * concrete model classes. Compatibility callers use specServeRun above. */
+export async function specRun(
+  binding: MlxSpeculativeBinding,
+  numDraftTokens: number,
+  promptIds: number[],
+  options: GenerateOptions & { stopSequences?: string[] },
+  onToken: OnToken,
+  services: SpecRunServices = {},
+): Promise<GenerateStats> {
+  return withRuntimeConfig(binding.runtime ?? runtimeConfig(), () => {
+    const run = () => specRunInner(binding, numDraftTokens, promptIds, options, onToken, services);
+    return binding.memory
+      ? withModelWiredLimit(binding.memory, () => withModelUsageFlush(binding.memory!, run))
+      : run();
+  });
+}
+
+async function specRunInner(
+  binding: MlxSpeculativeBinding,
+  numDraftTokens: number,
+  promptIds: number[],
+  options: GenerateOptions & { stopSequences?: string[] },
+  onToken: OnToken,
+  services: SpecRunServices = {},
 ): Promise<GenerateStats> {
   const maxTokens = options.maxTokens ?? 512;
-  const eos = options.eosTokenIds ?? model.config.eosTokenIds;
+  const maintainKv = createKvMaintenance(options);
+  const eos = options.eosTokenIds ?? [...binding.eosTokenIds];
   const sampler = makeSampler(options);
   const stepSampler = makeStepSampler(options, {
     tokenRepresentation: "number",
@@ -142,6 +133,8 @@ export async function specServeRun(
     sampler,
   });
   const gamma = Math.max(1, numDraftTokens);
+  const prefillChunk = options.prefillChunkSize ??
+    runtimeConfig().number("MLX_BUN_RD_PREFILL_CHUNK", PREFILL_CHUNK);
 
   // Allocated INSIDE the try below (2026-07-07 review): provider.open()
   // throws on a mismatched (target, drafter) pairing — as pre-try consts a
@@ -149,6 +142,7 @@ export async function specServeRun(
   // and turned a config error into a per-request 500. loadContext now
   // probe-opens the pairing at startup, so a throw here is belt+suspenders.
   let caches: Cache[] = [];
+  let prefixRetain: (() => void) | undefined;
   let source: import("./source").DraftSource | null = null;
   let tapLayers: number[] | undefined;
   // Target final hidden [1,1,H] at the anchor position — the assistant source
@@ -195,21 +189,73 @@ export async function specServeRun(
   const samplePos = async (logits1V: MlxArray, step: number): Promise<number> =>
     (await stepSampler.sample(logits1V, step)).token;
 
-  /** The [1,V] logits row at position `pos` of a hidden window — batched
-   *  lm-head is applied by the caller ONCE; this slices its output. */
-  const logitsRow = (logitsWindow: MlxArray, pos: number): MlxArray => {
-    const V = logitsWindow.shape[logitsWindow.shape.length - 1]!;
-    const sl = logitsWindow.slice([0, pos, 0], [1, pos + 1, V]);
-    const flat = ops.reshape(sl, [1, V]);
-    sl.dispose();
-    return flat;
+  const forwardTokens = async (tokens: number[], taps?: number[]) => {
+    const ids = ops.fromInt32(tokens, [1, tokens.length]);
+    try { return await binding.forward(ids, caches, taps); }
+    finally { ids.dispose(); }
   };
+  let failure: { error: unknown } | undefined;
 
   try {
-    caches = model.makeCache();
-    const src = provider.open({ sampler, target: { model, caches } });
+    options.signal?.throwIfAborted();
+    if (!Number.isSafeInteger(prefillChunk) || prefillChunk < 1)
+      throw new Error("speculative prefill chunk must be a positive integer");
+    assertMlxSpeculativeBinding(binding);
+    if (options.kvConfig?.length ||
+        (options.turboQuant && (options.quantizedKvStart ?? 0) !== 0) ||
+        (options.kvBits && ((options.kvBits !== 4 && options.kvBits !== 8) || options.quantizedKvStart !== 0)))
+      throw new Error("speculative KV requires uniform KV4/KV8 or TurboQuant start=0");
+    caches = binding.makeCache();
+    const src = binding.openDraft(sampler, caches);
     source = src;
     tapLayers = src.tapLayers;
+    const checkpoint = src.prefillMode === "full" ? src.checkpoint : undefined;
+    const prefixCache = checkpoint ? services.prefixCache : undefined;
+    const prefixNamespace = JSON.stringify({ method: checkpoint?.namespace,
+      adapters: services.cacheNamespace ?? options.adapters ?? [], kvBits: options.kvBits ?? null,
+      kvGroupSize: options.kvGroupSize ?? 64, quantizedKvStart: options.quantizedKvStart ?? null,
+      turboQuant: options.turboQuant ?? null });
+    let cached = 0;
+    const hit = prefixCache?.take(promptIds, prefixNamespace);
+    if (hit) {
+      try {
+        disposeResources(caches);
+        caches.splice(0, caches.length, ...hit.caches.splice(0));
+        prefixRetain = hit.retain;
+        hit.retain = undefined;
+        checkpoint!.restore(hit.tokens.length, hit.attachments![0]!);
+        cached = hit.tokens.length;
+      } finally {
+        disposeResources([...hit.caches, { dispose: () => disposeAttachments(hit.attachments) },
+          { dispose: () => hit.retain?.() }]);
+      }
+    }
+    stats.cachedTokens = cached;
+    const capture = (tokens: number[]) => {
+      let state: Cache[] = [];
+      let attachments: CheckpointAttachment[] = [];
+      try {
+        state = (services.cloneState ?? cloneKvCaches)(caches);
+        attachments.push(checkpoint!.capture(tokens.length));
+        prefixCache!.put(tokens, state, prefixNamespace, undefined, attachments);
+        state = []; attachments = [];
+      } finally {
+        disposeResources([...state, { dispose: () => disposeAttachments(attachments) }]);
+      }
+    };
+    const snapshotAt = prefixCache
+      ? Math.min(options.snapshotAt ?? promptIds.length, promptIds.length - 1) : -1;
+    const seedSource = async (end: number) => {
+      if (ctxParts.length === 1) prefillCtx = ctxParts[0]!;
+      else if (ctxParts.length > 1) {
+        prefillCtx = ops.concatAxis(ctxParts, 1);
+        for (const part of ctxParts) part.dispose();
+      }
+      ctxParts.length = 0;
+      const context = prefillCtx;
+      prefillCtx = null; // source.prefill owns context on both success and failure
+      await src.prefill(promptIds.slice(0, end), context ?? undefined);
+    };
     // ---- prefill (target, chunked; optionally tapped for DSpark), then seed
     // the source. Order: the source's prefill needs the tapped context that the
     // target prefill produces, so target-first (two-model's own draft prefill
@@ -240,77 +286,71 @@ export async function specServeRun(
     const t0 = performance.now();
     const tailSplit =
       src.prefillMode !== "full" &&
-      flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+      binding.prefillTailSplit;
     const oracleShape = tailSplit && promptIds.length > 1;
     if (oracleShape) {
       let pos = 0;
       const end = promptIds.length - 1; // last prompt token NEVER prefilled
       while (pos < end) {
-        const n = Math.min(PREFILL_CHUNK, end - pos);
+        options.signal?.throwIfAborted();
+        const n = Math.min(prefillChunk, end - pos);
         const chunk = promptIds.slice(pos, pos + n);
-        const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const { hidden: h, ctxML } = await forwardMaybeTap(
-          model,
-          ids,
-          caches,
-          tapLayers,
-        );
-        ids.dispose();
+        const { hidden: h, ctxML } = await forwardTokens(chunk, tapLayers);
         if (ctxML) ctxParts.push(ctxML);
-        if (pos + n >= end) {
-          const L = h.shape[1]!;
-          const H = h.shape[2]!;
-          anchorHidden = h.slice([0, L - 1, 0], [1, L, H]); // position len-2
-        }
-        h.dispose(); // logits never computed during the drain
+        try {
+          if (pos + n >= end) {
+            const L = h.shape[1]!;
+            const H = h.shape[2]!;
+            anchorHidden = h.slice([0, L - 1, 0], [1, L, H]); // position len-2
+          }
+        } finally { h.dispose(); } // no logits during the drain
+        maintainKv(caches);
         clearCache();
         pos += n;
+        if (pos < end) await new Promise<void>((resolve) => setImmediate(resolve));
       }
     } else {
       // Legacy shape: full chunked prefill (final chunk included), token0
       // sampled from the last position below.
-      for (let off = 0; off < promptIds.length; off += PREFILL_CHUNK) {
-        const chunk = promptIds.slice(off, off + PREFILL_CHUNK);
-        const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const { hidden: h, ctxML } = await forwardMaybeTap(
-          model,
-          ids,
-          caches,
-          tapLayers,
-        );
-        ids.dispose();
+      for (let off = cached; off < promptIds.length;) {
+        options.signal?.throwIfAborted();
+        const end = Math.min(off + prefillChunk, promptIds.length,
+          snapshotAt > off ? snapshotAt : Infinity);
+        const chunk = promptIds.slice(off, end);
+        const { hidden: h, ctxML } = await forwardTokens(chunk, tapLayers);
         if (ctxML) ctxParts.push(ctxML);
-        if (off + PREFILL_CHUNK >= promptIds.length) {
+        try { if (end >= promptIds.length) {
           const L = h.shape[1]!;
           const H = h.shape[2]!;
           anchorHidden = h.slice([0, L - 1, 0], [1, L, H]); // [1,1,H], retained
-          const lg = model.logitsFromHidden(anchorHidden);
-          const V = lg.shape[lg.shape.length - 1]!;
-          lastLogits = ops.reshape(lg, [1, V]);
-          lg.dispose();
+          const lg = binding.projectLogits(anchorHidden);
+          try {
+            const V = lg.shape[lg.shape.length - 1]!;
+            lastLogits = ops.reshape(lg, [1, V]);
+          } finally { lg.dispose(); }
+        } } finally { h.dispose(); }
+        maintainKv(caches);
+        if (end === snapshotAt && snapshotAt > cached) {
+          await seedSource(end);
+          capture(promptIds.slice(0, end));
         }
-        h.dispose();
         clearCache();
+        off = end;
+        if (off < promptIds.length)
+          await new Promise<void>((resolve) => setImmediate(resolve));
       }
     }
     extras.targetCalls++;
+    options.signal?.throwIfAborted();
     // Seed the source: two-model prefills its own draft cache; DSpark seeds
     // H_ctx from the tapped prompt context (ownership transfers to prefill);
     // the assistant is a no-op.
-    if (tapLayers) {
-      if (ctxParts.length === 1) prefillCtx = ctxParts[0]!;
-      else if (ctxParts.length > 1) {
-        prefillCtx = ops.concatAxis(ctxParts, 1); // [1,Lp,m*H]
-        for (const p of ctxParts) p.dispose();
-      }
-      ctxParts.length = 0; // parts consumed (transferred as prefillCtx or disposed)
-      await src.prefill(promptIds, prefillCtx ?? undefined); // takes ownership of prefillCtx
-      prefillCtx = null; // ownership transferred
-    } else {
-      await src.prefill(promptIds);
-    }
+    await seedSource(promptIds.length);
+    // Conversion replaces plain cache objects. Bind only after prefill has
+    // converted every attention layer; rollback must own the active objects.
+    const rollback = binding.bindRollback(caches);
     stats.prefillMs = performance.now() - t0;
-    stats.prefillTps = (promptIds.length / Math.max(stats.prefillMs, 1e-6)) * 1000;
+    stats.prefillTps = ((promptIds.length - cached) / Math.max(stats.prefillMs, 1e-6)) * 1000;
 
     // ---- first pending token ----
     // Oracle shape: pending = the UNPROCESSED last prompt token — never
@@ -326,11 +366,12 @@ export async function specServeRun(
       pending = await samplePos(lastLogits!, 0);
       lastLogits!.dispose();
       lastLogits = null;
+      stats.generatedTokens++;
       if (eos.includes(pending)) {
+        stats.finishReason = "stop";
         stats.decodeMs = performance.now() - tDecode;
         return stats;
       }
-      stats.generatedTokens++;
       if (
         (await onToken(pending)) === false ||
         stats.generatedTokens >= maxTokens ||
@@ -347,28 +388,24 @@ export async function specServeRun(
 
     // ---- rounds ----
     decode: while (stats.generatedTokens < maxTokens) {
+      options.signal?.throwIfAborted();
       if (!speculating) {
         // plain single-token continuation (post-ring-wrap fallback)
-        const ids = ops.fromInt32([pending], [1, 1]);
-        const { hidden: h } = await forwardMaybeTap(
-          model,
-          ids,
-          caches,
-          undefined,
-        );
-        ids.dispose();
-        const lg = model.logitsFromHidden(h);
-        h.dispose();
-        const V = lg.shape[lg.shape.length - 1]!;
-        roundRow = ops.reshape(lg, [1, V]);
-        lg.dispose();
+        const { hidden: h } = await forwardTokens([pending], undefined);
+        try {
+          const lg = binding.projectLogits(h);
+          try {
+            const V = lg.shape[lg.shape.length - 1]!;
+            roundRow = ops.reshape(lg, [1, V]);
+          } finally { lg.dispose(); }
+        } finally { h.dispose(); }
         extras.targetCalls++;
         const tok = await samplePos(roundRow, stats.generatedTokens);
         roundRow.dispose();
         roundRow = null;
         clearCache();
-        if (eos.includes(tok)) break;
         stats.generatedTokens++;
+        if (eos.includes(tok)) { stats.finishReason = "stop"; break; }
         pending = tok;
         if ((await onToken(tok)) === false) break;
         if (grammar?.isTerminated) break; // grammar closed — finish "stop"
@@ -387,11 +424,7 @@ export async function specServeRun(
       // spec-eligible through the round snapshot/replay contract instead:
       // specRoundBegin before the verify forward, then commit or a bit-exact
       // rollback(keep) after the accept walk.
-      const roundFits = caches.every((c) =>
-        "maxSize" in c && typeof (c as { maxSize?: number }).maxSize === "number"
-          ? c.offset + n + 1 < (c as { maxSize: number }).maxSize
-          : c.isTrimmable() || typeof c.specRoundBegin === "function",
-      );
+      const roundFits = rollback.canBegin(n);
       if (!roundFits) {
         console.warn(
           "spec: sliding window nearly wrapped — finishing without speculation",
@@ -432,21 +465,20 @@ export async function specServeRun(
       // the emitted position is the next round's anchor hidden.
       // Arm recurrent caches' spec round around the verify forward (no-op for
       // trimmable caches, which roll back via trim() below).
-      for (const c of caches) c.specRoundBegin?.();
+      rollback.begin(d);
       const vIds = ops.fromInt32([pending, ...drafts], [1, d + 1]);
-      const pinTarget = src.pinTargetKernelFamily === true &&
-        "setSpecKernelPinned" in model;
-      if (pinTarget) model.setSpecKernelPinned(true);
+      let pin: { close(): void } | undefined;
       try {
-        const ft = await forwardMaybeTap(model, vIds, caches, tapLayers);
+        pin = src.pinTargetKernelFamily ? binding.pinVerify?.() : undefined;
+        const ft = await binding.forward(vIds, caches, tapLayers);
         // Assign these before the lm-head call so the outer cleanup owns them
         // even if logits construction throws.
         vHidden = ft.hidden;
         vCtxML = ft.ctxML;
-        vLogits = model.logitsFromHidden(vHidden);
+        vLogits = binding.projectLogits(vHidden);
       } finally {
         vIds.dispose();
-        if (pinTarget) model.setSpecKernelPinned(false);
+        pin?.close();
       }
       extras.targetCalls++;
       // The lm-head remains one logical batched operation. Native GLM MTP's
@@ -457,39 +489,12 @@ export async function specServeRun(
       // accept while it reproduces the draft. Sampling is sequential because
       // the processor history (and, in Phase C, the grammar mask) at position
       // i depends on the tokens accepted at positions < i.
-      let kAccept = 0;
-      let correction: number | null = null; // target's token at first mismatch (or bonus)
-      let sawEos = false;
+      const [result] = await sampleSpeculativeRows([{
+        pending, step: stats.generatedTokens, remaining: maxTokens - stats.generatedTokens,
+        eosTokenIds: eos, sampling: stepSampler, grammarDone: () => grammar?.isTerminated ?? false,
+      }], [drafts], vLogits!);
+      const { accepted: kAccept, correction, sawEos, emitted, grammarDone } = result!.acceptance;
       let halted = false;
-      const emitted: number[] = [];
-      let grammarDone = false;
-      for (let i = 0; i <= d; i++) {
-        roundRow = logitsRow(vLogits!, i);
-        const tok = await samplePos(roundRow, stats.generatedTokens + emitted.length);
-        roundRow.dispose();
-        roundRow = null;
-        if (i < d && tok === drafts[i]) {
-          kAccept++;
-          // EOS is never content — even when it arrives as an ACCEPTED DRAFT
-          // (the correction/bonus branch below always excluded it; this
-          // branch pushed-then-broke, leaking the EOS through onToken —
-          // caught by the 2026-07-07 live oracle gate: streams were
-          // bit-identical, ours emitted one extra "content" token: <|eot_id|>).
-          if (eos.includes(tok)) { sawEos = true; break; }
-          emitted.push(tok);
-          // grammar termination mid-burst truncates the round — nothing may
-          // be sampled past a satisfied grammar (the all--inf guarantee).
-          if (grammar?.isTerminated) { grammarDone = true; break; }
-          if (stats.generatedTokens + emitted.length >= maxTokens) break;
-          continue;
-        }
-        // mismatch (target replaces the draft) or bonus (i === n)
-        correction = tok;
-        if (!eos.includes(tok)) emitted.push(tok);
-        else sawEos = true;
-        if (grammar?.isTerminated) grammarDone = true;
-        break;
-      }
       vLogits!.dispose();
       vLogits = null;
       extras.accepted += kAccept;
@@ -502,10 +507,18 @@ export async function specServeRun(
 
       // (d) emit the round's tokens through onToken, one at a time
       for (const tok of emitted) {
+        options.signal?.throwIfAborted();
         stats.generatedTokens++;
         if ((await onToken(tok)) === false) { halted = true; break; }
         if (stats.generatedTokens >= maxTokens) break;
       }
+      // Native generation and mlx-lm count the stopping EOS in usage while
+      // keeping it out of content. A callback or budget may stop earlier in
+      // this verified burst; that unused suffix must not enter usage.
+      if (sawEos && !halted && stats.generatedTokens < maxTokens) {
+        stats.generatedTokens++;
+        stats.finishReason = "stop";
+      }
       extras.tokensPerForward = extras.rounds > 0
         ? stats.generatedTokens / extras.rounds
         : 0;
@@ -526,14 +539,7 @@ export async function specServeRun(
         // drafts). Trimmable caches drop the rejected tail; recurrent caches
         // restore their pre-round snapshot and bit-exactly replay those
         // kAccept+1 kept tokens. On full accept the round just commits.
-        if (kAccept < d) {
-          for (const c of caches) {
-            if (c.specRoundRollback) c.specRoundRollback(kAccept + 1);
-            else c.trim(d - kAccept);
-          }
-        } else {
-          for (const c of caches) c.specRoundCommit?.();
-        }
+        rollback.resolve(kAccept);
         await src.commit(
           d,
           kAccept,
@@ -565,21 +571,15 @@ export async function specServeRun(
       ? (stats.generatedTokens / stats.decodeMs) * 1000
       : 0;
     return stats;
+  } catch (error) {
+    failure = { error };
+    throw error;
   } finally {
-    // Dispose whatever a mid-round/prefill throw left live (each is nulled the
-    // instant its normal disposal or ownership transfer ran, so no double-free).
-    anchorHidden?.dispose();
-    lastLogits?.dispose();
-    prefillCtx?.dispose();
-    for (const p of ctxParts) p.dispose();
-    vLogits?.dispose();
-    roundRow?.dispose();
-    vHidden?.dispose();
-    vCtxML?.dispose();
-    for (const c of caches) c.dispose();
-    source?.dispose();
-    stepSampler.dispose();
-    options.grammar?.dispose(); // Phase C will consume it; never leak either way
-    clearCache();
+    const resources = [anchorHidden, lastLogits, prefillCtx, ...ctxParts,
+      vLogits, roundRow, vHidden, vCtxML, ...caches, source, stepSampler,
+      options.grammar, { dispose: () => prefixRetain?.() }, { dispose: clearCache }]
+      .filter((resource): resource is DisposableResource => resource != null);
+    if (failure) cleanupFailure(failure.error, () => disposeResources(resources));
+    else disposeResources(resources);
   }
 }