mlx-bun 0.0.9 → 0.0.11

package/src/generate.ts

@@ -15,20 +15,34 @@ import {
 } from "./mlx/ffi";
 import * as ops from "./mlx/ops";
 import { CompiledDecode } from "./model/compiled-decode";
-import { Gemma4Model, KVCache, RotatingKVCache, type Cache } from "./model/gemma4";
+import { Gemma4Model, KVCache, RotatingKVCache, TurboQuantKVCache, type Cache } from "./model/gemma4";
+import { PagedKVCache } from "./model/paged-kv";
 import { DiffusionGemmaModel } from "./model/diffusion-gemma";
 import { diffusionGenerate } from "./diffusion/diffusion-generate";
 import type { RuntimeModel } from "./model/factory";
-import type { KvQuantSpec } from "./config";
+import type { KvQuantSpec, TurboQuantScheme } from "./config";
+import { flagOn } from "./flags";
 import {
   makeLogitsProcessors, makeSampler,
   type LogitsProcessorOptions, type SamplerOptions, toLogprobs,
 } from "./sampler";
+import type { GrammarController } from "./grammar";
 
 export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions {
   maxTokens?: number;
   eosTokenIds?: number[];
   prefillChunkSize?: number;
+  /** Fired once during prefill, at `snapshotAt` tokens (default: the full
+   *  prompt) and BEFORE any further KV is written — the one moment the
+   *  caches hold exactly that prefix (post-wrap rings can't be rewound
+   *  later). The server's prompt-boundary cache snapshot hangs here. */
+  onPrefillDone?: () => void;
+  /** Token count at which onPrefillDone fires (the STABLE cache boundary —
+   *  chat prompts end in a generation primer, e.g. 12B's thought-channel
+   *  tokens, that the NEXT turn's re-render does not contain; snapshotting
+   *  there would make the entry untrimmable-and-divergent). Prefill is
+   *  split at this index when it falls mid-prompt. */
+  snapshotAt?: number;
   /** Pre-warmed KV caches (e.g. from the prompt cache). cache[0].offset
    *  prompt tokens are treated as already prefilled; only the suffix is
    *  forwarded. Caller keeps ownership — generate() will not dispose. */
@@ -37,8 +51,16 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
    *  covering the whole prompt; prefilled in one shot (no chunking).
    *  Caller keeps ownership. */
   promptEmbeddings?: MlxArray;
-  /** bool [L] marking image tokens (bidirectional attention among them). */
+  /** bool [L] marking image tokens (bidirectional attention among them).
+   *  MUST be unset when the prompt contains any audio — audio prompts run
+   *  fully causal (audio-input-plan.md §3.3 Q1). */
   imageMask?: MlxArray;
+  /** bool [L] marking ALL multimodal soft tokens (image | audio) for
+   *  per-layer-input id zeroing (e2b/e4b), decoupled from imageMask so
+   *  audio-only prompts (no bidirectional mask) still zero their positions.
+   *  When unset, zeroing falls back to imageMask (the legacy vision shape).
+   *  Caller keeps ownership. */
+  multimodalMask?: MlxArray;
   /** Quantize full-attention KV caches to this many bits (4 or 8).
    *  Rotating (sliding-window) caches stay bf16 — they're window-capped
    *  and upstream rotating-cache quantization is NYI. */
@@ -52,6 +74,22 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
   /** Convert once a cache's offset reaches this (uniform-kvBits default
    *  5000 = mlx-lm; kvConfig default 0 = optiq serve). */
   quantizedKvStart?: number;
+  /** TurboQuant scheme (docs/design/turboquant-kv.md): rotation-based KV
+   *  quantization, a CLI-only runtime lever in the same class as uniform
+   *  kvBits (mutually exclusive with kvBits/kvConfig — maybeQuantizeKv
+   *  checks kvBits/kvConfig first, so set at most one). Full-attention
+   *  KVCache layers convert via TurboQuantKVCache.fromKVCache;
+   *  RotatingKVCache (sliding-window) layers stay bf16 in v1 — a one-time
+   *  warning names the limitation, never a throw. */
+  turboQuant?: TurboQuantScheme;
+  /** OPTIONAL paged KV storage (docs/design/paged-kv-cache.md): fresh
+   *  full-attention KVCache layers are replaced with PagedKVCache (block
+   *  pool + gather-to-contiguous) before prefill. v1 scope: serial batch=1
+   *  Gemma4-family, bf16 — mutually exclusive with kvBits/kvConfig/
+   *  turboQuant/draft/compiled decode (callers refuse the combos; the
+   *  cache swap itself only ever touches plain empty KVCache entries).
+   *  Values are gated bit-exact vs the plain path (tests/paged-kv*). */
+  pagedKv?: { blockSize?: number };
   /** Mounted LoRA adapter ids to apply (resolved/validated by
    *  AdapterManager.resolveSpec). Residuals sum in order. Set on the
    *  model's LoraState for exactly the duration of this generation —
@@ -70,6 +108,13 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
   /** Capture the top-k (token id, logprob) pairs per emitted token from the
    *  same distribution (mlx_lm.server `top_logprobs`). 0/unset = off. */
   topLogprobs?: number;
+  /** Grammar-constrained decoding (src/grammar.ts): a compiled GrammarController
+   *  that masks invalid tokens to -inf each step. L2-class (oMLX oracle). When
+   *  set, the decode loop takes a slightly different shape: it eager-reads the
+   *  token id (acceptToken needs a JS number, which the pipelined loop defers),
+   *  advances the matcher, and awaits the async mask precompute (which overlaps
+   *  the GPU forward). Non-grammar requests keep the fast pipelined loop. */
+  grammar?: GrammarController;
 }
 
 /** Per-token logprob info (only present when GenerateOptions requested it). */
@@ -82,14 +127,35 @@ export interface TokenLogprobs {
   top?: { id: number; logprob: number }[];
 }
 
+/** Materialize every cache's state at a prefill chunk boundary. Unlike
+ *  KVCache/QuantizedKVCache/RotatingKVCache, TurboQuantKVCache.state()
+ *  allocates FRESH trimmed slice views on every call (required by its
+ *  snapshotCache/cloneKvCaches callers in kv-store.ts, which already
+ *  dispose them) rather than returning its own live-owned arrays — so
+ *  this chokepoint must dispose that cache kind's state() output itself,
+ *  or the views leak (unreferenced past this call, only reclaimed by the
+ *  FinalizationRegistry backstop on GC of the tiny JS wrapper). */
+export function evalCacheState(cache: Cache[]): void {
+  const turboState = cache.flatMap((c) => (c instanceof TurboQuantKVCache ? c.state() : []));
+  try {
+    const liveState = cache.flatMap((c) => (c instanceof TurboQuantKVCache ? [] : c.state()));
+    ops.evalAll([...turboState, ...liveState]);
+  } finally {
+    // Disposed even when evalAll throws (Metal OOM / deferred graph error) —
+    // the whole reason this chokepoint exists is that these views are owned.
+    for (const a of turboState) a.dispose();
+  }
+}
+
 /** Port of mlx-lm maybe_quantize_kv_cache + BOTH halves of optiq serve's
  *  per-layer patched variant (incl. patch_rotating_to_quantized: rotating
  *  caches convert too — Phase 9):
  *  - per-layer bits/group_size selection (kvConfig overrides kvBits,
  *    matching optiq's --kv-config precedence; shipped kv_config.json
- *    files list full-attention layers only, so rotating quantization
- *    engages through uniform kvBits — like optiq --kv-bits — or a config
- *    that names sliding layers), and
+ *    files cover EVERY cache-owning layer, sliding ones included —
+ *    verified 12B 48/48, 26B 30/30, e4b 24/24 distinct caches — so
+ *    rotating quantization engages straight from the config; uniform
+ *    kvBits — like optiq --kv-bits — reaches them too), and
  *  - STREAMING conversion (optiq streaming_kv_quant / serve.py
  *    patched_maybe_quantize): eval each layer's quantized triples and
  *    clear the buffer pool before building the next layer's conversion.
@@ -99,7 +165,11 @@ export interface TokenLogprobs {
  *    24 GB Mac). Numerics untouched: same quantize math, only the eval
  *    ordering is forced (tests/kv-quant.test.ts, tests/rotating-kvq.test.ts). */
 export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void {
-  const { kvBits, kvConfig } = options;
+  const { kvBits, kvConfig, turboQuant } = options;
+  if (turboQuant) {
+    maybeTurboQuantizeKv(cache, turboQuant, options.quantizedKvStart ?? 0);
+    return;
+  }
   if (!kvBits && !kvConfig?.length) return;
   const start = options.quantizedKvStart ?? (kvConfig?.length ? 0 : 5000);
   const byLayer = kvConfig?.length
@@ -128,6 +198,70 @@ export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void
   }
 }
 
+/** Paged-KV conversion (docs/design/paged-kv-cache.md): swap each FRESH
+ *  plain full-attention KVCache for a PagedKVCache sized to hold
+ *  `capacityTokens` (prompt + maxTokens — known exactly at generate()
+ *  setup, so pool exhaustion is unreachable absent an accounting bug).
+ *  Same in-place-mutation shape as maybeQuantizeKv, but runs ONCE before
+ *  prefill: paging changes storage layout, not arithmetic, so there is no
+ *  "convert when populated" trigger. Sliding-window (RotatingKVCache)
+ *  layers keep today's scheme — mixed paged-full + rotating-sliding is
+ *  the supported v1 shape. Pre-warmed caches (offset > 0) skip conversion
+ *  entirely: the serve lane bypasses prompt-cache reuse for paged
+ *  requests, so this only guards library callers. */
+export function maybePageKv(
+  cache: Cache[], options: GenerateOptions, capacityTokens: number,
+): void {
+  if (!options.pagedKv) return;
+  if (cache.some((c) => c.offset > 0)) return;
+  const blockSize = options.pagedKv.blockSize ?? PagedKVCache.DEFAULT_BLOCK_SIZE;
+  for (let i = 0; i < cache.length; i++) {
+    if (cache[i] instanceof KVCache) {
+      cache[i]!.dispose(); // fresh (offset 0) — nothing stored yet
+      cache[i] = new PagedKVCache(capacityTokens, blockSize);
+    }
+  }
+}
+
+/** Emitted once per process: RotatingKVCache (sliding-window) layers are a
+ *  documented v1 non-goal (docs/design/turboquant-kv.md) — they stay bf16
+ *  rather than throwing, so mixed full-attention/sliding-window models (e.g.
+ *  Gemma) still serve correctly under --kv-quant turbo. */
+let warnedTurboRotating = false;
+
+/** TurboQuant conversion chokepoint (mirrors the uniform/config branch
+ *  above): only plain full-attention KVCache instances convert, via
+ *  TurboQuantKVCache.fromKVCache — RotatingKVCache stays bf16 (warn once,
+ *  never throw). Same offset===0 skip-empty-cache rule as the affine path. */
+function maybeTurboQuantizeKv(cache: Cache[], scheme: TurboQuantScheme, start: number): void {
+  for (let i = 0; i < cache.length; i++) {
+    const c = cache[i]!;
+    if (c instanceof RotatingKVCache) {
+      if (!warnedTurboRotating) {
+        warnedTurboRotating = true;
+        console.warn(
+          "[turbo-quant] sliding-window (RotatingKVCache) layers stay bf16 in v1 " +
+          "(full-attention only) — docs/design/turboquant-kv.md.",
+        );
+      }
+      continue;
+    }
+    if (!(c instanceof KVCache) || c.offset < start || c.offset === 0) continue;
+    const tq = TurboQuantKVCache.fromKVCache(c, scheme.kBits, scheme.vBits);
+    cache[i] = tq;
+    // state() allocates fresh trimmed slice views for this cache kind
+    // (see evalCacheState) — dispose after materializing (throw included),
+    // or they leak.
+    const state = tq.state();
+    try {
+      ops.evalAll(state);
+    } finally {
+      for (const a of state) a.dispose();
+    }
+    clearCache();
+  }
+}
+
 export interface GenerateStats {
   promptTokens: number;
   /** Prompt tokens skipped via a pre-warmed cache. */
@@ -141,6 +275,14 @@ export interface GenerateStats {
    *  decoded token that was forwarded, including a trailing EOS the
    *  pipeline forwarded before reading it). For PromptCache.put(). */
   cacheTokens: number[];
+  /** Speculative-decoding telemetry (serve --draft-model path only).
+   *  draftedByPos/acceptedByPos: per-draft-position round counts (index =
+   *  position within a round's block) — the Phase-1c per-position
+   *  acceptance signal. */
+  spec?: {
+    drafted: number; accepted: number; targetCalls: number;
+    draftedByPos?: number[]; acceptedByPos?: number[];
+  };
 }
 
 export interface GeneratedToken {
@@ -324,6 +466,11 @@ async function* generateInner(
     throw new Error(
       `pre-warmed cache (${cachedTokens} tokens) must be a strict prefix of the prompt (${promptTokens.length})`,
     );
+  // Paged KV (default off): swap fresh full-attention caches for paged
+  // storage BEFORE prefill. Capacity is exact — the deepest write lands at
+  // prompt + maxTokens − 1 (step maxTokens−1's forward), so this bound
+  // makes pool exhaustion an accounting-bug tripwire, not a request limit.
+  maybePageKv(cache, options, promptTokens.length + maxTokens);
   /** device-side token history (only maintained when processors need it) */
   let history: MlxArray | null = null;
 
@@ -367,6 +514,10 @@ async function* generateInner(
     const V = logits3d.shape[2]!;
     let logits = ops.reshape(logits3d, [1, V]);
     for (const p of processors) logits = disposing(logits, p(history, logits));
+    // Grammar mask (src/grammar.ts): apply AFTER the standard processors so the
+    // grammar has the final say on validity (a -inf'd token stays -inf through
+    // log-softmax). Off (no grammar) = zero cost — the branch is unset.
+    if (options.grammar) logits = disposing(logits, options.grammar.applyMask(logits));
     const lp = toLogprobs(logits);
     logits.dispose();
     const tok = sampler(lp, step);
@@ -444,21 +595,65 @@ async function* generateInner(
       if (needsTokenHistory)
         history = ops.fromInt32(promptTokens, [promptTokens.length]);
       // e2b/e4b need the spliced token ids to build per-layer inputs
-      // (image positions zeroed inside forwardEmbeddings).
+      // (multimodal soft-token positions zeroed inside forwardEmbeddings).
       const embedIds = ops.fromInt32(promptTokens, [1, promptTokens.length]);
       h0 = model.forwardEmbeddings(
         options.promptEmbeddings, cache, options.imageMask ?? null, embedIds,
+        options.multimodalMask ?? null,
       );
       embedIds.dispose();
     } else {
       let pos = cachedTokens;
-      while (promptTokens.length - pos > prefillChunkSize) {
-        const chunk = promptTokens.slice(pos, pos + prefillChunkSize);
+      // Stable-boundary snapshot: prefill up to snapshotAt first (plain
+      // chunks, no logits), fire the hook while the caches hold exactly
+      // that prefix, then continue with the remainder below. Splitting here
+      // changes chunk shapes only for requests that opted in — the same
+      // numerics class as a prompt-cache-hit continuation prefill.
+      const snapAt = options.onPrefillDone && options.snapshotAt !== undefined
+        ? Math.min(Math.max(options.snapshotAt, cachedTokens), promptTokens.length)
+        : promptTokens.length;
+      if (snapAt > pos && snapAt < promptTokens.length) {
+        while (pos < snapAt) {
+          const chunk = promptTokens.slice(pos, Math.min(pos + prefillChunkSize, snapAt));
+          const ids = ops.fromInt32(chunk, [1, chunk.length]);
+          const h = model.forwardHidden(ids, cache);
+          ids.dispose();
+          h.dispose();
+          evalCacheState(cache);
+          maybeQuantizeKv(cache, options);
+          clearCache();
+          pos += chunk.length;
+          // Macrotask yield between chunks (runtime-isolation.md Phase 1):
+          // each chunk is a synchronous multi-hundred-ms FFI eval; without
+          // this the event loop serves no I/O for the WHOLE prefill.
+          await new Promise<void>((r) => setImmediate(r));
+        }
+        options.onPrefillDone?.();
+      }
+      // Oracle prefill convention (mlx-lm 0.31.3 generate_step,
+      // generate.py:430-453): drain the prompt only to len-1 — chunks of
+      // min(prefillChunkSize, remaining-1) — then compute step-0 logits from
+      // a SEPARATE L=1 forward of the LAST prompt token (its server's
+      // batched engine shares the convention: insert_segments' forced final
+      // 1-token segment + GenerationBatch._step, generate.py:1645/1327).
+      // Forwarding the last token as the tail of an L=n GEMM instead is
+      // ulp-different in bf16 (qmm-vs-qmv + L-dependent SDPA dispatch) in
+      // BOTH the step-0 logits and that token's stored KV — near-tie greedy
+      // streams flip (2026-07-07 12B completion-probe ✗: first token flip at
+      // step 24). MLX_BUN_PREFILL_TAIL_SPLIT=0 restores the old
+      // full-final-chunk convention (A/B lever + kill switch).
+      const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+      const drainGate = tailSplit ? 1 : prefillChunkSize;
+      while (promptTokens.length - pos > drainGate) {
+        const n = tailSplit
+          ? Math.min(prefillChunkSize, promptTokens.length - pos - 1)
+          : prefillChunkSize;
+        const chunk = promptTokens.slice(pos, pos + n);
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
         const h = model.forwardHidden(ids, cache);
         ids.dispose();
         h.dispose(); // logits never computed for non-final chunks
-        ops.evalAll(cache.flatMap((c) => c.state()));
+        evalCacheState(cache);
         maybeQuantizeKv(cache, options);
         // mlx-lm _prefill clears the allocator cache after every chunk;
         // without this, prefill transients pile up in the buffer cache
@@ -467,16 +662,26 @@ async function* generateInner(
         // measured, scripts/decode-split.ts; the context-scaling decode
         // gap's main term).
         clearCache();
-        pos += prefillChunkSize;
+        pos += n;
+        // Macrotask yield between chunks (runtime-isolation.md Phase 1).
+        await new Promise<void>((r) => setImmediate(r));
       }
       if (needsTokenHistory) {
         history = ops.fromInt32(promptTokens, [promptTokens.length]);
       }
+      // Under tailSplit this is exactly [promptTokens[len-1]] — the L=1
+      // step-0 forward (uncompiled L=1 is bit-exact with compiled decode,
+      // tests/compiled-decode.test.ts). The last prompt token's KV enters
+      // the cache HERE, so downstream bookkeeping (forwarded[], cacheTokens,
+      // PromptCache.put alignment) is unchanged: after step 0 the caches
+      // cover exactly the prompt, as before.
       const lastChunk = promptTokens.slice(pos);
       const ids0 = ops.fromInt32(lastChunk, [1, lastChunk.length]);
       h0 = model.forwardHidden(ids0, cache);
       ids0.dispose();
     }
+    if (options.snapshotAt === undefined || options.snapshotAt >= promptTokens.length)
+      options.onPrefillDone?.();
     const [, L0, H] = h0.shape as [number, number, number];
     const hLast = h0.slice([0, L0 - 1, 0], [1, L0, H]);
     h0.dispose();
@@ -509,21 +714,110 @@ async function* generateInner(
     // growing attention windows change shape (= every step). Remove this
     // when upstream implements GatherQMM::output_shapes.
     let compiled =
-      process.env.MLX_BUN_COMPILED_DECODE !== "0" &&
+      flagOn("MLX_BUN_COMPILED_DECODE", true) &&
       !options.adapters?.length &&
+      // Paged caches can't compile (data-dependent block-list length —
+      // the shapeless-replay hazard; CompiledDecode.supports() also
+      // excludes them per step, this just skips the setup).
+      !options.pagedKv &&
       model.config.modelType.startsWith("gemma4") &&
       !model.config.text.enableMoeBlock
         ? CompiledDecode.for(model as Gemma4Model)
         : null;
     let stop = false;
+    /** Token id read eagerly at the top of the loop for grammar (reused for
+     *  the yield, avoiding a second readback). -1 when grammar is off — the
+     *  pipelined path keeps its deferred itemUint32 below. */
+    let grammarTok = -1;
+    // Jump-forward decoding (opt-in, MLX_BUN_GRAMMAR_JUMP=1; serial lane
+    // only — the batch lane's #stepGrammar doesn't jump yet): when the
+    // grammar forces a unique continuation, emit its retokenized ids without
+    // per-token forwards — ONE multi-token forward carries them into the KV
+    // (see GrammarController.jumpForward for the contract + the fidelity
+    // note on why this is opt-in). Excluded when logprobs are requested
+    // (jumped tokens are never sampled, so they'd have no logprobs rows).
+    const grammarJump =
+      options.grammar !== undefined &&
+      flagOn("MLX_BUN_GRAMMAR_JUMP", false) &&
+      !options.logprobs;
     while (!stop) {
       const cur = pending!;
       const curExtras = pendingExtras;
       pendingExtras = null;
+      // Grammar advance (src/grammar.ts): acceptToken needs the token id as a
+      // JS number, which the pipelined loop defers (it operates on device
+      // arrays). So grammar requests eager-read cur here, advance the matcher,
+      // and await the async mask precompute — which overlaps the GPU forward
+      // dispatched just below. This trades the readback/forward overlap for
+      // correctness (the mask for step n+1 must reflect token n). Non-grammar
+      // requests keep the fast pipelined loop untouched.
+      //
+      // F1 fix (batched-lane plan): always eager-read `grammarTok` when
+      // grammar is on, even at the max_tokens boundary — the emitted token
+      // reuses it unconditionally. The OLD code gated the READ on
+      // generated+1 < maxTokens but emitted grammarTok regardless, so the LAST
+      // iteration skipped the refresh and emitted a stale/garbage token (the
+      // previous step's token, or -1 when max_tokens=1) → truncated JSON ended
+      // on a corrupted token + cacheTokens recorded the wrong id. Now only
+      // accept()/ready() (which prepare the NEXT step's mask) are gated.
+      /** Forced ids to emit after cur this iteration (jump-forward), else null. */
+      let jumpEmit: number[] | null = null;
+      if (options.grammar) {
+        grammarTok = ops.itemUint32(cur);
+        if (generated + 1 < maxTokens) {
+          options.grammar.accept(grammarTok);
+          await options.grammar.ready();
+          if (grammarJump && !options.grammar.isTerminated) {
+            jumpEmit = options.grammar.jumpForward(maxTokens - (generated + 1));
+            // jumpForward advanced the matcher and fired the post-jump mask
+            // fill; it must be ready before this iteration's sampleStep.
+            if (jumpEmit) await options.grammar.ready();
+          }
+        }
+      }
       // build step n+1's graph from the *unread* pending token
       nextPending = null;
       nextExtras = null;
-      if (generated + 1 < maxTokens) {
+      // When the grammar has terminated (a complete valid JSON/schema
+      // accepted), there are no valid tokens left — skip building the next
+      // step so the sampler never sees an all--inf distribution.
+      if (jumpEmit) {
+        // JUMP iteration: one [1, 1+m] forward carries cur AND the forced ids
+        // into the KV (they are all committed content — jumpForward's
+        // contract); the next sampled token, if the budget and grammar allow
+        // one, comes from its last position. Compiled decode resumes on the
+        // following iteration (supports() re-checks the grown caches).
+        maybeQuantizeKv(cache, options);
+        pushHistory(cur);
+        if (needsTokenHistory) {
+          const jt = ops.fromInt32(jumpEmit, [jumpEmit.length]);
+          const prev = history!;
+          history = ops.concatAxis([prev, jt], 0);
+          prev.dispose();
+          jt.dispose();
+        }
+        const chunk = [grammarTok, ...jumpEmit];
+        const ids = ops.fromInt32(chunk, [1, chunk.length]);
+        const h = model.forwardHidden(ids, cache);
+        ids.dispose();
+        // Every chunk token's KV is in the cache regardless of what follows.
+        forwarded.push(...chunk);
+        const willGen = generated + 1 + jumpEmit.length;
+        if (willGen < maxTokens && !options.grammar!.isTerminated) {
+          const [, Lj, Hj] = h.shape as [number, number, number];
+          const hLast = h.slice([0, Lj - 1, 0], [1, Lj, Hj]);
+          h.dispose();
+          const logits = model.logitsFromHidden(hLast);
+          hLast.dispose();
+          const sn = sampleStep(logits, willGen);
+          nextPending = sn.tok;
+          nextExtras = sn.extras;
+          logits.dispose();
+          ops.asyncEvalAll([nextPending, ...extrasArrays(nextExtras)]);
+        } else {
+          h.dispose(); // burst ends the generation (max_tokens or grammar done)
+        }
+      } else if (generated + 1 < maxTokens && !options.grammar?.isTerminated) {
         maybeQuantizeKv(cache, options);
         pushHistory(cur);
         let logits: MlxArray | null = null;
@@ -558,7 +852,9 @@ async function* generateInner(
       }
 
       // sync-read step n's token while n+1 computes
-      const token = ops.itemUint32(cur);
+      // sync-read step n's token while n+1 computes (grammar already read it
+      // eagerly above — reuse to avoid a second GPU sync)
+      const token = options.grammar ? grammarTok : ops.itemUint32(cur);
       if (generated === 0) {
         // first token arrived: prompt clock stops, decode clock starts
         // (mlx-lm stream_generate's n==0 clock swap; the first token is
@@ -570,7 +866,8 @@ async function* generateInner(
       pending = null;
       generated++;
       // if a next-step graph was built, this token's KV entered the cache
-      if (nextPending !== null) forwarded.push(token);
+      // (jump iterations pushed the whole chunk already)
+      if (nextPending !== null && !jumpEmit) forwarded.push(token);
 
       if (eosTokenIds.includes(token)) {
         disposeExtras(curExtras);
@@ -587,6 +884,18 @@ async function* generateInner(
         // mlx-lm generate_step: clear_cache after token 0 (drops the
         // remaining prefill transients) and every 256 tokens after
         if ((generated - 1) % 256 === 0) clearCache();
+        // Jump-forward burst: the forced ids follow cur, one yield each (the
+        // consumer's stop-sequence matcher and detokenizer see the same
+        // one-at-a-time stream shape as always). Their KV is already in the
+        // cache (the chunk forward above); a consumer break mid-burst is
+        // safe — `forwarded` already reflects the cache exactly.
+        if (jumpEmit) {
+          for (const jt of jumpEmit) {
+            generated++;
+            yield { token: jt, index: generated - 1 };
+            if ((generated - 1) % 256 === 0) clearCache();
+          }
+        }
         if (nextPending === null) {
           stop = true;
         } else {
@@ -611,6 +920,11 @@ async function* generateInner(
       disposeExtras(nextExtras);
     }
     if (ownsCache) for (const c of cache) c.dispose();
+    // The grammar controller owns native WASM state (GrammarMatcher +
+    // TokenizerInfo cache) — dispose on every exit path (normal, early break,
+    // throw). TokenizerInfo itself is process-cached (vocab-structural), so
+    // only the per-request matcher/compiled are freed here.
+    options.grammar?.dispose();
     history?.dispose();
     if (!finished && !threw) {
       // forced early return (consumer break at a yield): still report