mlx-bun 0.3.0 → 0.4.0

package/src/generate.ts

@@ -10,6 +10,7 @@
 import { appendFileSync } from "node:fs";
 import { MlxArray, gpuStream } from "./mlx/array";
 import {
+  Dtype,
   activeMemory,
   clearCache,
   maxRecommendedWorkingSetSize,
@@ -18,8 +19,9 @@ import {
   synchronize,
 } from "./mlx/ffi";
 import * as ops from "./mlx/ops";
-import { KVCache, RotatingKVCache, TurboQuantKVCache, type Cache } from "./model/gemma4";
-import { PagedKVCache } from "./lab/paged-kv/paged-kv";
+import { KVCache, RotatingKVCache, type Cache } from "./model/gemma4";
+import { maybePageKv } from "./backends/mlx/request-state-policy";
+export { maybePageKv } from "./backends/mlx/request-state-policy";
 import { DiffusionGemmaModel } from "./model/diffusion-gemma";
 import { denoiseAsync } from "./diffusion/diffusion-generate";
 import { bindLegacyDenoisingModel, type MlxDenoisingBinding } from "./backends/mlx/diffusion";
@@ -28,6 +30,7 @@ import {
   assertMlxAutoregressiveBinding, bindLegacyAutoregressiveModel,
   type MlxAutoregressiveBinding, type MlxModelMemory, type MlxDecodeStep, type MlxTokenAppend,
 } from "./backends/mlx/autoregressive";
+import { createKvMaintenance } from "./backends/mlx/kv-maintenance";
 import { nextPrefillStep } from "./inference/prefill";
 import { evalCacheState, executeMlxPrefillStep } from "./backends/mlx/prefill";
 import { appendFillHidden } from "./backends/mlx/fill-append";
@@ -40,6 +43,7 @@ import { runtimeConfig, runtimeValue, withRuntimeConfig, type RuntimeConfig } fr
 import {
   disposeStepExtras,
   makeStepSampler,
+  readStepExtras,
   stepExtrasArrays,
   type LogitsProcessorOptions,
   type SamplerOptions,
@@ -57,8 +61,6 @@ import {
 } from "./fill/fill-session";
 
 export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions {
-  /** RAM budget for a model-owned paired target/draft prompt snapshot. */
-  speculativeCacheBytes?: number;
   /** Resolved host policy. Direct compatibility calls resolve from their
    * captured binding when this is absent. */
   decodePolicy?: Readonly<Pick<import("./contracts/execution").ResolvedExecution, "compiledDecode" | "grammarJump">>;
@@ -121,14 +123,13 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
    *  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. */
+   *  Full-attention and rotating caches both convert when populated. */
   kvBits?: number;
   kvGroupSize?: number;
   /** Per-layer mixed-precision KV from kv_config.json (config.kvQuant).
    *  Overrides kvBits, like optiq serve's --kv-config. layerIdx indexes
-   *  the cache list (== layer index for the donor prefix); entries for
-   *  rotating/sliding caches are skipped until Phase 9. */
+   *  the cache list (== layer index for the donor prefix), including
+   *  rotating/sliding caches. */
   kvConfig?: KvQuantSpec[];
   /** Convert once a cache's offset reaches this (uniform-kvBits default
    *  5000 = mlx-lm; kvConfig default 0 = optiq serve). */
@@ -136,14 +137,14 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
   /** TurboQuant scheme (docs/design/turboquant.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
+   *  checks turboQuant 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/kv-cache.md): fresh
    *  full-attention KVCache layers are replaced with PagedKVCache (block
-   *  pool + gather-to-contiguous) before prefill. v1 scope: serial batch=1
+   *  pool + gather-to-contiguous) before prefill. Scope: B1/B>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).
@@ -230,88 +231,12 @@ export function shouldUseFill(
     options.turboQuant === undefined;
 }
 
-/** 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 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.
- *    Lazily batching every layer's toQuantized into one eval pins ALL
- *    layers' bf16 K/V as graph inputs alongside ALL quantized outputs —
- *    the exact transient optiq's fix kills (16.35 → 7.60 GB at 32k on a
- *    24 GB Mac). Numerics untouched: same quantize math, only the eval
- *    ordering is forced (tests/parity/kv-quant.test.ts, tests/parity/rotating-kvq.test.ts). */
+/** Compatibility entry point for callers that perform a single conversion.
+ * Execution sessions compose createKvMaintenance once and reuse it. */
 export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void {
-  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
-    ? new Map(kvConfig.map((e) => [e.layerIdx, e]))
-    : null;
-  for (let i = 0; i < cache.length; i++) {
-    const c = cache[i]!;
-    if (!(c instanceof KVCache || c instanceof RotatingKVCache) || c.offset < start) continue;
-    // OptiQ's mixed-KV hook skips empty caches: the first prompt prefill
-    // runs bf16, then the populated cache is quantized before decode.
-    // Converting empty caches at start=0 makes prefill itself quantized
-    // and diverges from the oracle path.
-    if (c.offset === 0) continue;
-    if (byLayer) {
-      const e = byLayer.get(i);
-      if (!e) continue;
-      cache[i] = c.toQuantized(e.groupSize, e.bits);
-    } else {
-      cache[i] = c.toQuantized(options.kvGroupSize ?? 64, kvBits!);
-    }
-    // Streaming half: materialize THIS layer's conversion now so its bf16
-    // source (already released by toQuantized) frees before the next layer
-    // converts — the transient stays ~one layer, not the whole cache.
-    ops.evalAll(cache[i]!.state());
-    clearCache();
-  }
-}
-
-/** Paged-KV conversion (docs/design/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);
-    }
-  }
+  createKvMaintenance(options)(cache);
 }
 
-/** Emitted once per process: RotatingKVCache (sliding-window) layers are a
- *  documented v1 non-goal (docs/design/turboquant.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;
-
 /** Emitted once per process: token fast-forwarding skips models with
  *  sliding-window layers in v1 (see the gate in generateInner). */
 let warnedFillRotating = false;
@@ -330,39 +255,6 @@ type RewindableCache = {
 };
 const rewindable = (c: unknown): RewindableCache => c as RewindableCache;
 
-/** 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.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;
   /** An observed terminal cause takes precedence over the token-count limit. */
@@ -455,9 +347,11 @@ export class Generation implements AsyncIterable<GeneratedToken> {
 // Scope semantics match the reference: set → generate → synchronize →
 // restore; nothing stays pinned between generations. Re-entrant: only
 // the outermost wiring scope touches the limit.
-// macOS 26.6 reports a 24.96 GiB recommended set on this 24 GB machine.
-// The old 0.75 fraction therefore stopped wiring the 13-16 GiB Qwen/GLM
-// models even though they page heavily without an explicit wired limit.
+// macOS 26.6 reports a 26.8e9 B (24.96 GiB, 0.78 × RAM) recommended set on
+// a 32 GB M1 Max (2026-09-08; the earlier "24 GB machine" attribution was
+// wrong — 24.96 GiB exceeds a 24 GB box's RAM). At 0.75 × that the old
+// fraction stopped wiring the 13-16 GiB Qwen/GLM models even though they
+// page heavily without an explicit wired limit.
 // Keep smaller 8-9 GiB models unwired while covering the large-model class.
 const WIRE_THRESHOLD = 0.5;
 let wiredScopeDepth = 0;
@@ -719,6 +613,7 @@ async function* generateInner(
     throw error;
   }
   const graph = binding.graph;
+  const maintainKv = createKvMaintenance(options);
 
   // logprobs capture (mlx_lm.server semantics — see GenerateOptions.logprobs).
   // Everything below is gated: when neither flag is set, no extra ops, evals,
@@ -797,25 +692,6 @@ async function* generateInner(
     (c) => c.isTrimmable() || typeof rewindable(c).specRoundRollback === "function",
   );
   closeBatchSetup?.();
-  /** Device-side logprob capture for one step, computed from the SAME lp the
-   *  sampler saw (post-processors, pre-truncation) — read back lazily with the
-   *  token so decode pipelining is preserved. */
-  /** Read extras back to JS (forces eval — they were async-dispatched with the
-   *  token) without appending casts behind the next decode step. */
-  const readExtras = (e: StepExtras | null): TokenLogprobs | undefined => {
-    if (!e) return undefined;
-    const out: TokenLogprobs = {};
-    if (e.sel) out.logprob = e.sel.toFloat32Host()[0]!;
-    if (e.topIdx && e.topVals) {
-      const ids = e.topIdx.toIntTokens();
-      const vals = e.topVals.toFloat32Host();
-      out.top = Array.from(ids, (id, i) => ({ id, logprob: vals[i]! })).sort(
-        (a, b) => b.logprob - a.logprob,
-      );
-    }
-    disposeStepExtras(e);
-    return out;
-  };
 
   // logits [1,1,V] → sampled token array [1] (+ optional logprob capture,
   // all on-device)
@@ -848,6 +724,7 @@ async function* generateInner(
   let finished = false;
   let threw = false;
   let executionError: unknown;
+  let earlyUnforwardedToken: number | null = null;
   let closeTokenZero: (() => void) | undefined;
   const makeStats = (): GenerateStats => ({
     promptTokens: options.originalPromptTokens ?? promptTokens.length,
@@ -948,7 +825,7 @@ async function* generateInner(
     nextExtras = null;
     // No-op under the fill exclusions (kv quant refuses fill); kept so the
     // boundary stays correct if those exclusions ever loosen.
-    maybeQuantizeKv(cache, options);
+    maintainKv(cache);
 
     let checkpointMs = 0;
     let roundOpen = false;
@@ -1061,7 +938,10 @@ async function* generateInner(
       // pending token was sampled from that exact state before the snapshot,
       // so resume starts directly at the decode loop without another forward.
       if (needsTokenHistory) stepSampler.seedHistory(promptTokens);
-      pending = ops.fromInt32([options.initialPendingToken!], [1]);
+      // Device samplers return uint32. Preserve that signature on restore so
+      // compiled decode does not trace a different signed-token graph first.
+      using restoredToken = ops.fromInt32([options.initialPendingToken!], [1]);
+      pending = restoredToken.astype(Dtype.uint32);
       ops.asyncEvalAll([pending]);
       prefillMs = 0;
       closePrefill?.();
@@ -1072,7 +952,7 @@ async function* generateInner(
       mechanism: diagnostics.mechanism ?? "serial",
       boundary: "initial",
     });
-    maybeQuantizeKv(cache, options);
+    maintainKv(cache);
     closeInitialKv?.();
     tPrefill = performance.now();
     let h0: MlxArray;
@@ -1096,7 +976,7 @@ async function* generateInner(
         : undefined;
       const tailSplit = runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true);
       const forward = graph.forwardHidden.bind(graph);
-      const maintain = () => maybeQuantizeKv(cache, options);
+      const maintain = () => maintainKv(cache);
       while (true) {
         options.signal?.throwIfAborted();
         const step = nextPrefillStep({ length: promptTokens.length, position: pos,
@@ -1186,14 +1066,16 @@ async function* generateInner(
           closeTokenZero?.();
           closeTokenZero = undefined;
           tDecode = performance.now();
-          const logprobs = readExtras(curExtras);
+          const logprobs = readStepExtras(curExtras);
           curExtras = null;
           pendingExtras = null;
           // Count before yielding so an early return reports this token. Its
           // KV is still absent; forwarded describes the prompt-only cache.
           generated++;
           yieldedEarly = true;
+          earlyUnforwardedToken = firstToken;
           yield { token: firstToken, index: stepIndex, ...(logprobs ? { logprobs } : {}) };
+          earlyUnforwardedToken = null;
           options.signal?.throwIfAborted();
         }
       }
@@ -1242,7 +1124,7 @@ async function* generateInner(
         // 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);
+        maintainKv(cache);
         pushHistory(cur);
         stepSampler.commitNumbers(jumpEmit);
         const chunk = [grammarTok, ...jumpEmit];
@@ -1267,7 +1149,7 @@ async function* generateInner(
           h.dispose(); // burst ends the generation (max_tokens or grammar done)
         }
       } else if (stepIndex + 1 < maxTokens && !options.grammar?.isTerminated) {
-        maybeQuantizeKv(cache, options);
+        maintainKv(cache);
         pushHistory(cur);
         let logits: MlxArray | null = null;
         let evalWith: MlxArray[] = [];
@@ -1347,7 +1229,7 @@ async function* generateInner(
         }
         // readExtras before the yield: if the consumer breaks at this yield,
         // the extras are already read and disposed.
-        const logprobs = readExtras(curExtras);
+        const logprobs = readStepExtras(curExtras);
         pendingExtras = null;
         if (!yieldedEarly) yield { token, index: generated - 1, ...(logprobs ? { logprobs } : {}) };
         options.signal?.throwIfAborted();
@@ -1412,31 +1294,53 @@ async function* generateInner(
     executionError = e;
     throw e;
   } finally {
-    if (!finished) {
-      pending?.dispose();
-      nextPending?.dispose();
-      disposeStepExtras(pendingExtras);
-      disposeStepExtras(nextExtras);
-    }
     try {
-      const closing = decoder?.close();
-      if (closing) await closing;
+      // An early consumer return leaves token zero outside the retained KV.
+      // Finish that one M=1 forward before handing caller-owned caches back,
+      // matching the ordinary pipeline's boundary. Otherwise a later prefix
+      // hit folds the token into an M>1 prefill and changes native reductions.
+      // Aborted requests and caches owned by this run need no retained state.
+      if (!ownsCache && !threw && !options.signal?.aborted && earlyUnforwardedToken !== null) {
+        maintainKv(cache);
+        const ids = ops.reshape(pending!, [1, 1]);
+        try {
+          const hidden = await graph.forwardHidden(ids, cache);
+          try { evalCacheState(cache); }
+          finally { hidden.dispose(); }
+        } finally { ids.dispose(); }
+        forwarded.push(earlyUnforwardedToken);
+      }
     } catch (error) {
-      if (threw) throw new AggregateError([executionError, error],
-        "AR execution and decoder cleanup failed", { cause: executionError });
+      threw = true;
+      executionError = error;
       throw error;
     } finally {
-      if (ownsCache) for (const c of cache) c.dispose();
-      // TokenizerInfo is process-cached; this request owns the matcher.
-      options.grammar?.dispose();
-      stepSampler.dispose();
-    }
-    if (!finished && !threw) {
-      // forced early return (consumer break at a yield): still report
-      // stats — `forwarded` only lists tokens whose KV actually entered
-      // the cache, so cacheTokens stays exact for PromptCache.put().
-      decodeMs = performance.now() - tDecode;
-      return makeStats();
+      if (!finished) {
+        pending?.dispose();
+        nextPending?.dispose();
+        disposeStepExtras(pendingExtras);
+        disposeStepExtras(nextExtras);
+      }
+      try {
+        const closing = decoder?.close();
+        if (closing) await closing;
+      } catch (error) {
+        if (threw) throw new AggregateError([executionError, error],
+          "AR execution and decoder cleanup failed", { cause: executionError });
+        throw error;
+      } finally {
+        if (ownsCache) for (const c of cache) c.dispose();
+        // TokenizerInfo is process-cached; this request owns the matcher.
+        options.grammar?.dispose();
+        stepSampler.dispose();
+      }
+      if (!finished && !threw) {
+        // forced early return (consumer break at a yield): still report
+        // stats — `forwarded` only lists tokens whose KV actually entered
+        // the cache, so cacheTokens stays exact for PromptCache.put().
+        decodeMs = performance.now() - tDecode;
+        return makeStats();
+      }
     }
   }
 }

package/src/inference/draft-acceptance.ts

@@ -0,0 +1,41 @@
+/** One request's exact-token speculative acceptance walk. The backend owns
+ * logits and sampling; this state owns only the accepted prefix and stop facts.
+ * A group advances each live walk independently, regardless of batch size. */
+export class DraftAcceptance {
+  position = 0;
+  accepted = 0;
+  correction: number | null = null;
+  readonly emitted: number[] = [];
+  sawEos = false;
+  grammarDone = false;
+  done = false;
+
+  constructor(
+    readonly drafts: readonly number[],
+    readonly remaining: number,
+    readonly eosTokenIds: readonly number[],
+  ) {}
+
+  /** Supply one target sample after its sampler has updated request history
+   * and grammar. Once done, this request needs no further target samples. */
+  accept(token: number, grammarDone = false): void {
+    const matches = this.position < this.drafts.length && token === this.drafts[this.position];
+    this.position++;
+    if (matches) {
+      this.accepted++;
+      if (this.eosTokenIds.includes(token)) {
+        this.sawEos = this.done = true;
+        return;
+      }
+      this.emitted.push(token);
+      this.grammarDone = grammarDone;
+      this.done = grammarDone || this.emitted.length >= this.remaining;
+      return;
+    }
+    this.correction = token;
+    this.sawEos = this.eosTokenIds.includes(token);
+    if (!this.sawEos) this.emitted.push(token);
+    this.grammarDone = grammarDone;
+    this.done = true;
+  }
+}

package/src/inference/draft-output.ts

@@ -0,0 +1,49 @@
+import type { Cancellation, CancelReason, GenerationOutput, TokenLogprobs } from "../contracts/generation";
+import type { DraftAcceptance } from "./draft-acceptance";
+
+export interface DraftOutputRow {
+  readonly acceptance: DraftAcceptance;
+  readonly logprobs?: readonly (TokenLogprobs | undefined)[];
+  readonly output: Pick<GenerationOutput, "commit">;
+  readonly cancellation?: Cancellation;
+}
+
+/** Counts apply to this round. Accepted counts describe processed proposed
+ * inputs, excluding an unprocessed correction/bonus and any unpublished EOS.
+ * Failed/cancelled requests must be discarded, never checkpointed. */
+export type DraftOutputResult = { readonly generated: number; readonly accepted: number } & (
+  | { readonly kind: "continue"; readonly pending: number }
+  | { readonly kind: "stop" | "length" }
+  | { readonly kind: "cancelled"; readonly reason: CancelReason }
+  | { readonly kind: "failed"; readonly error: unknown }
+);
+
+/** Deliver verified content through the output interface before resolving
+ * retained state. A stopped or failed consumer affects only its own request.
+ * Token-sized publications preserve callbacks that stop inside a draft burst. */
+export async function deliverDraftOutputs(rows: readonly DraftOutputRow[]): Promise<DraftOutputResult[]> {
+  const results: DraftOutputResult[] = [];
+  for (const row of rows) results.push(await deliver(row));
+  return results;
+}
+
+async function deliver(row: DraftOutputRow): Promise<DraftOutputResult> {
+  const { acceptance } = row;
+  let generated = 0;
+  const retained = () => ({ generated, accepted: Math.min(generated, acceptance.accepted) });
+  try {
+    for (const token of acceptance.emitted) {
+      const reason = row.cancellation?.reason;
+      if (reason) return { kind: "cancelled", reason, generated, accepted: 0 };
+      const metadata = row.logprobs?.[generated];
+      const more = await row.output.commit([token], metadata ? [metadata] : undefined);
+      generated++;
+      if (more === false) return { kind: "stop", ...retained() };
+    }
+    const reason = row.cancellation?.reason;
+    if (reason) return { kind: "cancelled", reason, generated, accepted: 0 };
+    if (acceptance.sawEos || acceptance.grammarDone) return { kind: "stop", ...retained() };
+    if (generated >= acceptance.remaining) return { kind: "length", ...retained() };
+    return { kind: "continue", pending: acceptance.correction!, ...retained() };
+  } catch (error) { return { kind: "failed", error, generated, accepted: 0 }; }
+}

package/src/inference/rollback.ts

@@ -1,7 +1,7 @@
 /** A complete transaction over every participating target state layer.
  * A failed begin/resolve invalidates the run; its state must be discarded. */
-export interface SpeculativeTransaction {
+export interface SpeculativeTransaction<Acceptance = number> {
   canBegin(drafts: number): boolean;
   begin(drafts: number): void;
-  resolve(accepted: number): void;
+  resolve(accepted: Acceptance): void;
 }

package/src/kv-scheme.ts

@@ -86,6 +86,42 @@ export function kvBytesAt(
     value.linearStateBytes;
 }
 
+/** Head dims MLX 0.32.2's fused multi-query SDPA kernel serves for L > 8
+ *  queries on M1–M4 GPUs (mlx/backend/metal/scaled_dot_product_attention.cpp
+ *  `use_fallback`: 192 and 256 always fall back off-NAX; `has_fused_kernel`
+ *  lists 64/72/80/96/128 for the full kernel). Every other head dim runs
+ *  mlx/fast.cpp's composed fallback, which materializes the bf16 scores
+ *  tensor [heads, L, S] plus a bool mask per full-attention layer. The
+ *  decode kernel (L ≤ 8) fuses 256, so this term is a prefill-only cost. */
+export const FUSED_SDPA_HEAD_DIMS: ReadonlySet<number> = new Set([64, 72, 80, 96, 128]);
+const SDPA_SCORES_BYTES = 2; // bf16 scores in q's dtype
+const SDPA_FUSED_MAX_QUERIES = 8;
+
+/** Peak bytes of the SDPA-fallback scores tensor live during one prefill
+ *  chunk of `chunk` queries against a `ctx`-token context: the last chunk
+ *  sees the longest keys, and the lazy graph keeps one layer's scores live
+ *  at a time. Zero when every attention layer's head dim fuses. Sliding
+ *  layers see at most `window + L` keys. Admission-only (fit and the
+ *  in-flight memory guard); nothing here changes what is computed. */
+export function sdpaFallbackBytes(config: ModelConfig, chunk: number, ctx: number): number {
+  const text = config.text;
+  const heads = text.numAttentionHeads || 0;
+  const queries = Math.min(chunk, ctx);
+  if (heads <= 0 || queries <= SDPA_FUSED_MAX_QUERIES) return 0;
+  let peak = 0;
+  const layers = text.numHiddenLayers || text.layerTypes.length;
+  for (let layer = 0; layer < layers; layer++) {
+    const type = text.layerTypes[layer] ?? "full_attention";
+    if (type === "linear_attention") continue;
+    const sliding = type === "sliding_attention" && text.slidingWindow > 0;
+    const headDim = sliding ? text.headDim : text.globalHeadDim;
+    if (FUSED_SDPA_HEAD_DIMS.has(headDim)) continue;
+    const keys = sliding ? Math.min(ctx, text.slidingWindow + queries) : ctx;
+    peak = Math.max(peak, heads * queries * keys * SDPA_SCORES_BYTES);
+  }
+  return peak;
+}
+
 export class KvScheme {
   readonly kind: KvSchemeKind;
   readonly options: Readonly<ResolvedKvSchemeOptions>;
@@ -112,10 +148,17 @@ export class KvScheme {
   get cacheKey(): string {
     if (this.kind === "turbo") {
       const turbo = this.options.turboQuant!;
-      return `turbo-k${turbo.kBits}v${turbo.vBits}`;
+      const start = this.options.quantizedKvStart ?? 0;
+      return `turbo-k${turbo.kBits}v${turbo.vBits}${start > 0 ? `-start${start}` : ""}`;
+    }
+    if (this.kind === "affine-uniform") {
+      const start = this.options.quantizedKvStart ?? 5000, group = this.options.kvGroupSize ?? 64;
+      return `kv${this.options.kvBits}${group !== 64 ? `-g${group}` : ""}${start > 0 ? `-start${start}` : ""}`;
+    }
+    if (this.kind === "affine-config") {
+      const start = this.options.quantizedKvStart ?? 0;
+      return `config${start > 0 ? `-start${start}` : ""}`;
     }
-    if (this.kind === "affine-uniform") return `kv${this.options.kvBits}`;
-    if (this.kind === "affine-config") return "config";
     return "bf16";
   }
 
@@ -154,14 +197,24 @@ export class KvScheme {
     return this.kind === "turbo" ? {} : this.generationOptions;
   }
 
-  /** The current batch engine accepts only per-layer affine schemes whose
-   * named cache instances can perform the required conversion. The capability
-   * probe is mandatory for quantized schemes so config shape alone can never
-   * authorize placement. */
-  batchable(config: ModelConfig, canConvert?: (layerIdx: number) => boolean): boolean {
+  /** Affine layouts convert during solo prefill. TQ also has a transitional
+   * row layout that preserves each request's delayed conversion boundary. */
+  batchable(config: ModelConfig, canConvert?: (layerIdx: number) => boolean,
+    cacheLayerCount = config.text.numHiddenLayers, capabilities: { delayedAffine?: boolean } = {}): boolean {
     if (this.kind === "bf16") return true;
-    if (this.kind !== "affine-config") return false;
+    if (this.kind !== "affine-config" && this.kind !== "affine-uniform" && this.kind !== "turbo") return false;
     if (!canConvert) return false;
+    const start = this.options.quantizedKvStart ?? (this.kind === "affine-uniform" ? 5000 : 0);
+    if (start > 0 && this.kind !== "turbo" && !capabilities.delayedAffine) return false;
+    if (this.kind === "affine-uniform" || this.kind === "turbo") {
+      // Some models share the donor prefix's KV in later layers. Probe the
+      // actual cache list, just as uniform conversion does.
+      for (let layer = 0; layer < cacheLayerCount; layer++) {
+        if (config.text.layerTypes[layer] === "linear_attention") continue;
+        if (!canConvert(layer)) return false;
+      }
+      return true;
+    }
     return this.options.kvConfig!.every((entry) => {
       if (entry.layerIdx < 0 || entry.layerIdx >= config.text.numHiddenLayers) return false;
       return (config.text.layerTypes[entry.layerIdx] ?? "full_attention") !== "linear_attention" &&
@@ -175,11 +228,15 @@ export function resolveKvScheme(input: {
   turboQuant?: TurboQuantScheme;
   config?: readonly KvQuantSpec[] | null;
   missingConfig?: "bf16" | "error";
+  quantizedKvStart?: number;
 }): KvScheme {
+  const start = input.quantizedKvStart;
+  if (start !== undefined && (!Number.isSafeInteger(start) || start < 0))
+    throw new Error("quantizedKvStart must be a nonnegative integer");
   if (input.turboQuant) {
     return new KvScheme("turbo", {
       turboQuant: input.turboQuant,
-      quantizedKvStart: 0,
+      quantizedKvStart: start ?? 0,
     });
   }
   if (input.override === "off" || input.override === undefined)
@@ -190,10 +247,11 @@ export function resolveKvScheme(input: {
         throw new Error("model has no kv_config.json (--kv-quant config)");
       return new KvScheme("bf16", {});
     }
-    return new KvScheme("affine-config", { kvConfig: [...input.config] });
+    return new KvScheme("affine-config", { kvConfig: [...input.config],
+      ...(start === undefined ? {} : { quantizedKvStart: start }) });
   }
   return new KvScheme("affine-uniform", {
     kvBits: input.override,
-    quantizedKvStart: 0,
+    quantizedKvStart: start ?? 0,
   });
 }