mlx-bun 0.0.13 → 0.3.0

package/src/generate.ts

@@ -4,34 +4,74 @@
 // - decode pipelining via mx.async_eval: step n+1's graph is built and
 //   dispatched before step n's token is read back, so the GPU never idles
 //   on the JS round-trip
+// - opt-in early token-zero yield reduces latency before that pipeline starts
 // - sampling stays on-device; only the chosen token id crosses to JS
 
+import { appendFileSync } from "node:fs";
 import { MlxArray, gpuStream } from "./mlx/array";
 import {
+  activeMemory,
   clearCache,
   maxRecommendedWorkingSetSize,
+  peakMemory,
   setWiredLimit,
   synchronize,
 } from "./mlx/ffi";
 import * as ops from "./mlx/ops";
-import { CompiledDecode } from "./model/compiled-decode";
-import { Gemma4Model, KVCache, RotatingKVCache, TurboQuantKVCache, type Cache } from "./model/gemma4";
-import { PagedKVCache } from "./model/paged-kv";
+import { KVCache, RotatingKVCache, TurboQuantKVCache, type Cache } from "./model/gemma4";
+import { PagedKVCache } from "./lab/paged-kv/paged-kv";
 import { DiffusionGemmaModel } from "./model/diffusion-gemma";
-import { diffusionGenerate } from "./diffusion/diffusion-generate";
+import { denoiseAsync } from "./diffusion/diffusion-generate";
+import { bindLegacyDenoisingModel, type MlxDenoisingBinding } from "./backends/mlx/diffusion";
 import type { RuntimeModel } from "./model/factory";
+import {
+  assertMlxAutoregressiveBinding, bindLegacyAutoregressiveModel,
+  type MlxAutoregressiveBinding, type MlxModelMemory, type MlxDecodeStep, type MlxTokenAppend,
+} from "./backends/mlx/autoregressive";
+import { nextPrefillStep } from "./inference/prefill";
+import { evalCacheState, executeMlxPrefillStep } from "./backends/mlx/prefill";
+import { appendFillHidden } from "./backends/mlx/fill-append";
+export { evalCacheState } from "./backends/mlx/prefill";
+import type { TokenLogprobs } from "./contracts/generation";
+export type { TokenLogprobs } from "./contracts/generation";
+import type { PromptResponseTrace } from "./serve/prompt-response-trace";
 import type { KvQuantSpec, TurboQuantScheme } from "./config";
-import { flagOn } from "./flags";
+import { runtimeConfig, runtimeValue, withRuntimeConfig, type RuntimeConfig } from "./runtime-config";
 import {
-  makeLogitsProcessors, makeSampler,
-  type LogitsProcessorOptions, type SamplerOptions, toLogprobs,
+  disposeStepExtras,
+  makeStepSampler,
+  stepExtrasArrays,
+  type LogitsProcessorOptions,
+  type SamplerOptions,
+  type StepExtras,
 } from "./sampler";
 import type { GrammarController } from "./grammar";
+import {
+  fillTraceEnabled,
+  fillTracePath,
+  type FillTraceRecord,
+  resolveFillMode,
+  type FillSession,
+  type FillStats,
+  type Proposal,
+} 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">>;
+  /** Cooperatively cancel at prefill/decode boundaries; pending native work
+   *  finishes before owned resources are released. */
+  signal?: AbortSignal;
   maxTokens?: number;
   eosTokenIds?: number[];
   prefillChunkSize?: number;
+  /** Whether the request supplied its seed explicitly. Durable generation
+   *  checkpoints use this to distinguish a reproducible user policy from the
+   *  server-generated seed that must be recovered from checkpoint metadata. */
+  seedWasExplicit?: boolean;
   /** 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
@@ -43,6 +83,25 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
    *  there would make the entry untrimmable-and-divergent). Prefill is
    *  split at this index when it falls mid-prompt. */
   snapshotAt?: number;
+  /** Resume an interrupted autoregressive generation from a durable KV
+   *  checkpoint. `cache` must cover every token in `promptTokens` exactly;
+   *  `initialPendingToken` is the already-sampled next token, so no cache
+   *  rewind is required (important for recurrent/untrimmable layers). */
+  initialPendingToken?: number;
+  /** Number of completion tokens already emitted before this resumed run. */
+  initialGeneratedTokens?: number;
+  /** Original request prompt length used for protocol usage accounting when
+   *  `promptTokens` also contains replayed completion tokens. */
+  originalPromptTokens?: number;
+  /** Periodic, safe decode-boundary checkpoint hook. The cache covers
+   *  `cacheTokens`; `pendingToken` has been sampled but not emitted. */
+  checkpointEveryTokens?: number;
+  onDecodeCheckpoint?: (state: {
+    cacheTokens: number[];
+    caches: Cache[];
+    generatedTokens: number;
+    pendingToken: number;
+  }) => void | Promise<void>;
   /** 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. */
@@ -53,7 +112,7 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
   promptEmbeddings?: MlxArray;
   /** 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). */
+   *  fully causal (docs/design/generic-model-support.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
@@ -74,7 +133,7 @@ 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
+  /** 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
@@ -82,7 +141,7 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
    *  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
+  /** 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
    *  Gemma4-family, bf16 — mutually exclusive with kvBits/kvConfig/
@@ -115,45 +174,60 @@ export interface GenerateOptions extends SamplerOptions, LogitsProcessorOptions
    *  advances the matcher, and awaits the async mask precompute (which overlaps
    *  the GPU forward). Non-grammar requests keep the fast pipelined loop. */
   grammar?: GrammarController;
+  /** Token fast-forwarding (src/fill/, docs/design/speculative-decoding.md
+   *  "Token fast-forwarding"): a per-request table of DETERMINED token spans
+   *  (compiled from the request's `tools` + the chat template). When the
+   *  stream enters one, the engine appends the whole span itself with ONE
+   *  multi-token forward and resumes sampling after it.
+   *
+   *  This is context extension, NOT speculation: nothing is drafted, verified,
+   *  or rolled back, and no comparison is made against what the model would
+   *  have produced. Injection bypasses the sampler, which is a documented
+   *  behavior-policy deviation at temperature > 0 — hence opt-in
+   *  (MLX_BUN_FILL=strict) and the composition refusals in shouldUseFill.
+   *  Serial lane only: the batch scheduler and the spec loop never read it. */
+  fill?: FillSession;
 }
 
 export function shouldUseGrammarJump(
   options: Pick<GenerateOptions, "grammar" | "logprobs" | "topLogprobs">,
+  runtime: RuntimeConfig = runtimeConfig(),
 ): boolean {
   return options.grammar !== undefined &&
-    flagOn("MLX_BUN_GRAMMAR_JUMP", false) &&
+    runtime.flag("MLX_BUN_GRAMMAR_JUMP", false) &&
     !options.logprobs &&
     !(options.topLogprobs && options.topLogprobs > 0);
 }
 
-/** Per-token logprob info (only present when GenerateOptions requested it). */
-export interface TokenLogprobs {
-  /** logprob of the emitted token (requires options.logprobs). */
-  logprob?: number;
-  /** top-k (id, logprob) pairs, sorted by logprob descending (requires
-   *  options.topLogprobs > 0). mlx-lm leaves argpartition order unspecified;
-   *  sorting is the deterministic reading of the same set. */
-  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();
-  }
+/** Composition gate for token fast-forwarding, enforced inside the engine
+ *  (the serve layer refuses more shapes it alone can see — a user-fixed seed,
+ *  a mounted draft model, continuous placement: src/serve/request-prep.ts,
+ *  src/serve/chat-stage.ts). Throws when MLX_BUN_FILL names the unbuilt echo
+ *  tier — an operator who asked for K3c must not get K3b silently.
+ *   - grammar: forced-token production is the grammar's job; two producers in
+ *     one iteration is forbidden (asserted in the loop).
+ *   - logprobs/top_logprobs: injected tokens are never sampled, so they have
+ *     no distribution row (same rule as shouldUseGrammarJump).
+ *   - promptEmbeddings: vision/audio prompts (and mRoPE) — untested shape.
+ *   - kvBits/kvConfig/turboQuant: post-conversion multi-token append is
+ *     L-generic but unvalidated; refuse in v1 rather than pay it silently. */
+export function shouldUseFill(
+  options: Pick<
+    GenerateOptions,
+    "fill" | "grammar" | "logprobs" | "topLogprobs" | "promptEmbeddings"
+    | "kvBits" | "kvConfig" | "turboQuant"
+  >,
+  runtime: RuntimeConfig = runtimeConfig(),
+): boolean {
+  if (!options.fill) return false;
+  if (resolveFillMode(runtime.value("MLX_BUN_FILL") ?? "off") === "off") return false;
+  return options.grammar === undefined &&
+    !options.logprobs &&
+    !(options.topLogprobs && options.topLogprobs > 0) &&
+    options.promptEmbeddings === undefined &&
+    !options.kvBits &&
+    !options.kvConfig?.length &&
+    options.turboQuant === undefined;
 }
 
 /** Port of mlx-lm maybe_quantize_kv_cache + BOTH halves of optiq serve's
@@ -172,7 +246,7 @@ export function evalCacheState(cache: Cache[]): void {
  *    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/kv-quant.test.ts, tests/rotating-kvq.test.ts). */
+ *    ordering is forced (tests/parity/kv-quant.test.ts, tests/parity/rotating-kvq.test.ts). */
 export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void {
   const { kvBits, kvConfig, turboQuant } = options;
   if (turboQuant) {
@@ -207,7 +281,7 @@ export function maybeQuantizeKv(cache: Cache[], options: GenerateOptions): void
   }
 }
 
-/** Paged-KV conversion (docs/design/paged-kv-cache.md): swap each FRESH
+/** 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).
@@ -233,11 +307,29 @@ export function maybePageKv(
 }
 
 /** Emitted once per process: RotatingKVCache (sliding-window) layers are a
- *  documented v1 non-goal (docs/design/turboquant-kv.md) — they stay bf16
+ *  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;
+
+/** The rewind surface a verify-policy fill needs. `Cache` declares the
+ *  spec-round trio as optional members, but the cache list is a UNION with
+ *  GLM's MLACache (which declares none of them), so the union has no such
+ *  property to read. One narrowing helper keeps the call sites honest — every
+ *  use is still optional-chained. */
+type RewindableCache = {
+  isTrimmable(): boolean;
+  trim(n: number, bypass?: boolean): void;
+  specRoundBegin?(): void;
+  specRoundCommit?(): void;
+  specRoundRollback?(keep: number): void;
+};
+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,
@@ -250,7 +342,7 @@ function maybeTurboQuantizeKv(cache: Cache[], scheme: TurboQuantScheme, start: n
         warnedTurboRotating = true;
         console.warn(
           "[turbo-quant] sliding-window (RotatingKVCache) layers stay bf16 in v1 " +
-          "(full-attention only) — docs/design/turboquant-kv.md.",
+          "(full-attention only) — docs/design/turboquant.md.",
         );
       }
       continue;
@@ -273,6 +365,8 @@ function maybeTurboQuantizeKv(cache: Cache[], scheme: TurboQuantScheme, start: n
 
 export interface GenerateStats {
   promptTokens: number;
+  /** An observed terminal cause takes precedence over the token-count limit. */
+  finishReason?: "stop" | "length";
   /** Prompt tokens skipped via a pre-warmed cache. */
   cachedTokens: number;
   generatedTokens: number;
@@ -294,6 +388,14 @@ export interface GenerateStats {
     rejected?: number; rounds?: number; acceptanceLengths?: number[];
     tokensPerForward?: number; forwardsSaved?: number;
   };
+  /** Token fast-forwarding telemetry (serial lane, MLX_BUN_FILL=strict).
+   *  Present only when a fill table was actually armed for this generation. */
+  fill?: FillStats;
+}
+
+export interface GenerateDiagnostics {
+  trace?: PromptResponseTrace;
+  mechanism?: "serial" | "continuous";
 }
 
 export interface GeneratedToken {
@@ -305,16 +407,23 @@ export interface GeneratedToken {
 
 export class Generation implements AsyncIterable<GeneratedToken> {
   stats: GenerateStats | null = null;
-  readonly #iter: AsyncGenerator<GeneratedToken, GenerateStats>;
+  readonly #next: () => Promise<IteratorResult<GeneratedToken, GenerateStats>>;
+  readonly #return: () => Promise<IteratorResult<GeneratedToken, GenerateStats>>;
 
-  constructor(iter: AsyncGenerator<GeneratedToken, GenerateStats>) {
-    this.#iter = iter;
+  constructor(iter: AsyncGenerator<GeneratedToken, GenerateStats>, runtime = runtimeConfig()) {
+    // Async generators resume in the caller's context on every next/return.
+    // Bind both operations so consumer awaits and early close cannot change
+    // the runtime snapshot seen by kernels or their cleanup.
+    const next = iter.next.bind(iter);
+    const close = () => iter.return(undefined as unknown as GenerateStats);
+    this.#next = () => withRuntimeConfig(runtime, next);
+    this.#return = () => withRuntimeConfig(runtime, close);
   }
 
   async *[Symbol.asyncIterator](): AsyncGenerator<GeneratedToken> {
     try {
       while (true) {
-        const r = await this.#iter.next();
+        const r = await this.#next();
         if (r.done) {
           this.stats = r.value;
           return;
@@ -327,7 +436,7 @@ export class Generation implements AsyncIterable<GeneratedToken> {
       // disposal, wired/adapter scopes) and capture the stats its
       // early-return path still reports.
       if (this.stats === null) {
-        const r = await this.#iter.return(undefined as unknown as GenerateStats);
+        const r = await this.#return();
         if (r.done && r.value) this.stats = r.value;
       }
     }
@@ -346,18 +455,15 @@ 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.
-const WIRE_THRESHOLD = 0.75;
+// 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.
+// Keep smaller 8-9 GiB models unwired while covering the large-model class.
+const WIRE_THRESHOLD = 0.5;
 let wiredScopeDepth = 0;
 let wiredOldLimit = 0;
 
-type WiredModelMemory = {
-  readonly weightsBytes: number;
-  readonly expertRuntime?: {
-    readonly plan: { readonly plannedBytes: number };
-    flushUsage?: () => void;
-    finishUsage?: () => Promise<void>;
-  } | null;
-};
+type WiredModelMemory = MlxModelMemory;
 
 /** Bytes whose MLX graph must stay resident while a model executes. */
 export function wiredWorkingSetBytes(model: WiredModelMemory): number {
@@ -370,7 +476,7 @@ export function wiredWorkingSetBytes(model: WiredModelMemory): number {
 export function modelNeedsWiredLimit(
   model: WiredModelMemory,
   recommendedBytes = maxRecommendedWorkingSetSize(),
-  force = process.env.MLX_BUN_FORCE_WIRE === "1",
+  force = runtimeValue("MLX_BUN_FORCE_WIRE") === "1",
 ): boolean {
   return force ||
     wiredWorkingSetBytes(model) > WIRE_THRESHOLD * recommendedBytes;
@@ -387,24 +493,77 @@ function exitWiredScope(): void {
   }
 }
 
+/** Hold the process wired limit while one model owns GPU execution. */
+export function acquireModelWiredLimit(model: WiredModelMemory): () => void {
+  if (!modelNeedsWiredLimit(model)) return () => {};
+  enterWiredScope();
+  let released = false;
+  return () => {
+    if (released) return;
+    released = true;
+    exitWiredScope();
+  };
+}
+
 export function generate(
   model: RuntimeModel,
   promptTokens: number[],
   options: GenerateOptions = {},
+  diagnostics: GenerateDiagnostics = {},
 ): Generation {
+  return bindGeneration(model)(promptTokens, options, diagnostics);
+}
+
+/** Resolve the compatibility model once at host construction. Every request
+ * borrows the same binding and immutable runtime snapshot. */
+export function bindGeneration(model: RuntimeModel): (
+  promptTokens: number[], options?: GenerateOptions, diagnostics?: GenerateDiagnostics,
+) => Generation {
   // DiffusionGemma is non-autoregressive: route to the denoising engine instead
   // of the AR decode loop. Same Generation/GenerateStats contract so the CLI and
   // server stream it through the existing token machinery.
-  let inner =
-    model instanceof DiffusionGemmaModel
-      ? generateDiffusionInner(model, promptTokens, options)
-      : generateInner(model, promptTokens, options);
-  if (options.adapters?.length) inner = adapterScoped(model, options.adapters, inner);
-  const memoryModel: WiredModelMemory = model;
-  if (memoryModel.expertRuntime?.flushUsage)
-    inner = usageScoped(memoryModel, inner);
-  const wire = modelNeedsWiredLimit(model);
-  return new Generation(wire ? wiredScoped(inner) : inner);
+  if (!(model instanceof DiffusionGemmaModel)) {
+    const binding = bindLegacyAutoregressiveModel(model);
+    return (prompt, options, diagnostics) => generateAutoregressive(binding, prompt, options, diagnostics);
+  }
+  const binding = bindLegacyDenoisingModel(model);
+  return (prompt, options) => generateDenoising(binding, prompt, options);
+}
+
+/** Execute a supplied backend binding through the same AR loop and resource
+ * scopes as generate(). Callers serialize shared model/adapter access. Replacing
+ * this binding replaces graph, cache construction, media, and optional decode
+ * together; it never borrows another model's compiled execution implicitly. */
+export function generateAutoregressive(
+  binding: MlxAutoregressiveBinding,
+  promptTokens: number[],
+  options: GenerateOptions = {},
+  diagnostics: GenerateDiagnostics = {},
+): Generation {
+  const runtime = binding.runtime ?? runtimeConfig();
+  let inner = generateInner(binding, promptTokens, options, diagnostics, runtime);
+  if (options.adapters?.length && binding.adapters) {
+    inner = adapterScoped({ loraState: binding.adapters }, options.adapters, inner);
+  }
+  if (binding.memory.expertRuntime?.finishUsage || binding.memory.expertRuntime?.flushUsage)
+    inner = usageScoped(binding.memory, inner);
+  return new Generation(modelNeedsWiredLimit(binding.memory, undefined,
+    runtime.value("MLX_BUN_FORCE_WIRE") === "1") ? wiredScoped(inner) : inner, runtime);
+}
+
+/** Denoising bindings use their own state and feedback graph. Callers hold
+ * the native runtime's exclusive generation lease, including its global RNG. */
+export function generateDenoising<State>(
+  binding: MlxDenoisingBinding<State>, promptTokens: number[], options: GenerateOptions = {},
+): Generation {
+  let inner = generateDiffusionInner(binding, promptTokens, options);
+  if (options.adapters?.length && binding.adapters)
+    inner = adapterScoped({ loraState: binding.adapters }, options.adapters, inner);
+  if (binding.memory.expertRuntime?.finishUsage || binding.memory.expertRuntime?.flushUsage)
+    inner = usageScoped(binding.memory, inner);
+  const runtime = binding.runtime ?? runtimeConfig();
+  return new Generation(modelNeedsWiredLimit(binding.memory, undefined,
+    runtime.value("MLX_BUN_FORCE_WIRE") === "1") ? wiredScoped(inner) : inner, runtime);
 }
 
 /** Non-autoregressive diffusion generation, adapted to the AR Generation
@@ -412,11 +571,12 @@ export function generate(
  *  streams the emitted tokens. v1: greedy (temperature 0, confidence-threshold
  *  sampler — the OptiQ default); per-block intra-stream + temperature>0
  *  (categorical) are follow-ups. */
-async function* generateDiffusionInner(
-  model: DiffusionGemmaModel,
+async function* generateDiffusionInner<State>(
+  binding: MlxDenoisingBinding<State>,
   promptTokens: number[],
   options: GenerateOptions,
 ): AsyncGenerator<GeneratedToken, GenerateStats> {
+  options.signal?.throwIfAborted();
   const t0 = performance.now();
   const maxTokens = options.maxTokens ?? 256;
   // A fresh random canvas seed per request unless the caller pins one.
@@ -426,40 +586,52 @@ async function* generateDiffusionInner(
       : BigInt(Math.floor(Math.random() * 0x7fffffff));
   // The shipped checkpoint's tokenizer stops on {1, 106}; union any caller eos.
   const eos = [...new Set([1, 106, ...(options.eosTokenIds ?? [])])];
-  const result = diffusionGenerate(model, promptTokens, {
+  const result = await denoiseAsync(binding.graph, promptTokens, {
     maxTokens,
     sampler: "confidence-threshold",
     temperature: 0,
     eosTokenIds: eos,
     seed,
     visionPixels: options.visionPixels,
-  });
+  }, options.signal);
   const decodeMs = performance.now() - t0;
   let index = 0;
-  for (const token of result.tokens) yield { token, index: index++ };
-  return {
-    promptTokens: promptTokens.length,
-    cachedTokens: 0,
-    generatedTokens: result.tokens.length,
-    prefillTps: 0,
-    decodeTps: result.tokens.length / Math.max(decodeMs / 1000, 1e-9),
-    prefillMs: 0,
-    decodeMs,
-    cacheTokens: [],
-  };
+  let failure: { error: unknown } | undefined;
+  try {
+    for (const token of result.tokens) {
+      options.signal?.throwIfAborted();
+      yield { token, index: index++ };
+    }
+  } catch (error) { failure = { error }; }
+  finally {
+    // A stop parser may close delivery before the finished canvas is exhausted.
+    // Settle the emitted count on return(), without suppressing cancellation/errors.
+    if (failure) throw failure.error;
+    return {
+      promptTokens: promptTokens.length,
+      cachedTokens: 0,
+      generatedTokens: index,
+      prefillTps: 0,
+      decodeTps: index / Math.max(decodeMs / 1000, 1e-9),
+      prefillMs: 0,
+      decodeMs,
+      cacheTokens: [],
+    };
+  }
 }
 
 /** Hold the model's active-adapter list for exactly this generation. */
 async function* adapterScoped(
-  model: RuntimeModel,
+  model: { loraState: { active: string[] } },
   adapters: string[],
   inner: AsyncGenerator<GeneratedToken, GenerateStats>,
 ): AsyncGenerator<GeneratedToken, GenerateStats> {
+  const previous = model.loraState.active;
   model.loraState.active = adapters;
   try {
     return yield* inner;
   } finally {
-    model.loraState.active = [];
+    model.loraState.active = previous;
   }
 }
 
@@ -496,12 +668,11 @@ export async function withModelWiredLimit<T>(
   model: WiredModelMemory,
   run: () => Promise<T>,
 ): Promise<T> {
-  if (!modelNeedsWiredLimit(model)) return run();
-  enterWiredScope();
+  const release = acquireModelWiredLimit(model);
   try {
     return await run();
   } finally {
-    exitWiredScope();
+    release();
   }
 }
 
@@ -520,37 +691,34 @@ export async function withModelUsageFlush<T>(
   }
 }
 
-/** Forward through either the ordinary synchronous model surface or a
- * streamed-expert model whose layer execution must cross async I/O. Keeping
- * this choice inside the shared generator preserves the complete generation
- * contract (sampling, grammar, logprobs, prompt cache, and cancellation) for
- * both execution styles. */
-async function forwardHiddenForGeneration(
-  model: RuntimeModel,
-  ids: MlxArray,
-  cache: Cache[],
-): Promise<MlxArray> {
-  const asyncModel = model as RuntimeModel & {
-    forwardHiddenAsync?: (ids: MlxArray, cache: Cache[]) => Promise<MlxArray>;
-  };
-  return typeof asyncModel.forwardHiddenAsync === "function"
-    ? await asyncModel.forwardHiddenAsync(ids, cache)
-    : model.forwardHidden(ids, cache);
-}
-
 async function* generateInner(
-  model: RuntimeModel,
+  binding: MlxAutoregressiveBinding,
   promptTokens: number[],
   options: GenerateOptions,
+  diagnostics: GenerateDiagnostics,
+  runtime: RuntimeConfig,
 ): AsyncGenerator<GeneratedToken, GenerateStats> {
+  if (options.signal?.aborted) {
+    options.grammar?.dispose();
+    options.signal.throwIfAborted();
+  }
   const {
     maxTokens = 512,
-    eosTokenIds = model.config.eosTokenIds,
+    eosTokenIds = binding.eosTokenIds,
     prefillChunkSize = 2048,
   } = options;
-  const sampler = makeSampler(options);
-  const processors = makeLogitsProcessors(options);
-  const needsTokenHistory = processors.length > 0;
+
+  try {
+    assertMlxAutoregressiveBinding(binding);
+    if (options.promptEmbeddings && !binding.forwardEmbeddings)
+      throw new Error("AR binding does not support prompt embeddings");
+    if (options.adapters?.length && !binding.adapters)
+      throw new Error("AR binding does not support adapters");
+  } catch (error) {
+    options.grammar?.dispose();
+    throw error;
+  }
+  const graph = binding.graph;
 
   // logprobs capture (mlx_lm.server semantics — see GenerateOptions.logprobs).
   // Everything below is gated: when neither flag is set, no extra ops, evals,
@@ -558,36 +726,80 @@ async function* generateInner(
   const captureSelLp = options.logprobs === true;
   const captureTopK =
     options.topLogprobs && options.topLogprobs > 0 ? options.topLogprobs : 0;
-  const capture = captureSelLp || captureTopK > 0;
+  const stepSampler = makeStepSampler(options, {
+    tokenRepresentation: "device",
+    grammarWait: "external",
+    historyUpdate: "manual",
+    captureSelectedLogprob: captureSelLp,
+    captureTopLogprobs: captureTopK,
+  });
+  const needsTokenHistory = stepSampler.needsHistory;
 
+  const closePrefill = diagnostics.trace?.begin("prefill.total", {
+    mechanism: diagnostics.mechanism ?? "serial",
+    promptTokens: promptTokens.length,
+    cachedTokens: options.cache?.[0]?.offset ?? 0,
+  });
+  const closeBatchSetup = diagnostics.trace?.begin("prefill.batch_setup", {
+    mechanism: diagnostics.mechanism ?? "serial",
+  });
   const ownsCache = !options.cache;
-  const cache = options.cache ?? model.makeCache();
+  const resuming = options.initialPendingToken !== undefined;
+  let cache: Cache[] = options.cache ?? [];
+  try {
+    if (ownsCache) cache = binding.makeCache();
+    if (!cache.length) throw new Error("AR binding returned an empty cache");
+    const cached = cache[0]!.offset;
+    if ((!resuming && cached >= promptTokens.length) ||
+        (resuming && cached !== promptTokens.length))
+      throw new Error(
+        resuming
+          ? `resumed cache (${cached} tokens) must exactly cover the resume prefix (${promptTokens.length})`
+          : `pre-warmed cache (${cached} tokens) must be a strict prefix of the prompt (${promptTokens.length})`,
+      );
+    // Replace fresh full-attention caches before any forward. The deepest
+    // write is prompt + maxTokens - 1, so capacity covers every decode step.
+    maybePageKv(cache, options, promptTokens.length + maxTokens);
+  } catch (error) {
+    if (ownsCache) for (const state of cache) state.dispose();
+    stepSampler.dispose();
+    options.grammar?.dispose();
+    closeBatchSetup?.();
+    closePrefill?.();
+    throw error;
+  }
   const cachedTokens = cache[0]!.offset;
-  if (cachedTokens >= promptTokens.length)
-    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;
-
+  // Token fast-forwarding gate (see shouldUseFill). RotatingKVCache (sliding
+  // window) layers append multi-token writes through #updateConcat — O(window)
+  // per append rather than the O(L) a plain ring pays — so v1 warns and skips
+  // the whole feature for those models rather than paying it silently.
+  let fillOn = shouldUseFill(options, runtime);
+  if (fillOn && cache.some((c) => c instanceof RotatingKVCache)) {
+    fillOn = false;
+    if (!warnedFillRotating) {
+      warnedFillRotating = true;
+      console.warn(
+        "[fill] sliding-window (RotatingKVCache) layers skip token " +
+        "fast-forwarding in v1 (multi-token append is O(window) there) — " +
+        "docs/design/speculative-decoding.md.",
+      );
+    }
+  }
+  const fillTrace = fillOn && fillTraceEnabled();
+  const fillTraceFile = fillOn ? fillTracePath() : null;
+  // Verify-policy proposals (echo tier) need the rejected tail rewound. That
+  // is the SAME contract the spec lane's rounds use: trimmable caches drop the
+  // tail; recurrent caches (SSMCache — gated-DeltaNet state, untrimmable)
+  // restore a pre-round snapshot and bit-exactly replay the accepted prefix.
+  // A model whose caches can do neither still gets assert-policy fills; verify
+  // proposals are dropped and counted (stats.verifyUnsupported).
+  const verifyCapable = fillOn && cache.every(
+    (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. */
-  interface StepExtras {
-    sel: MlxArray | null; // [1,1] emitted token's logprob
-    topIdx: MlxArray | null; // [1,k] top-k token ids
-    topVals: MlxArray | null; // [1,k] their logprobs
-  }
-  const extrasArrays = (e: StepExtras | null): MlxArray[] =>
-    e ? [e.sel, e.topIdx, e.topVals].filter((a): a is MlxArray => a !== null) : [];
-  const disposeExtras = (e: StepExtras | null): void => {
-    if (e) for (const a of [e.sel, e.topIdx, e.topVals]) a?.dispose();
-  };
   /** 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 => {
@@ -601,7 +813,7 @@ async function* generateInner(
         (a, b) => b.logprob - a.logprob,
       );
     }
-    disposeExtras(e);
+    disposeStepExtras(e);
     return out;
   };
 
@@ -611,52 +823,11 @@ async function* generateInner(
     logits3d: MlxArray,
     step: number,
   ): { tok: MlxArray; extras: StepExtras | null } => {
-    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);
-    let extras: StepExtras | null = null;
-    if (capture) {
-      let sel: MlxArray | null = null;
-      let topIdx: MlxArray | null = null;
-      let topVals: MlxArray | null = null;
-      if (captureSelLp) {
-        const idx = ops.reshape(tok, [1, 1]);
-        sel = ops.takeAlongAxis(lp, idx, -1);
-        idx.dispose();
-      }
-      if (captureTopK > 0) {
-        // mlx-lm _format_top_logprobs: argpartition(-logprobs, kth=k-1)[:k]
-        const k = Math.min(captureTopK, V);
-        const negLp = ops.neg(lp);
-        const part = ops.argpartitionAxis(negLp, k - 1, -1);
-        const idxView = part.slice([0, 0], [1, k]);
-        topIdx = ops.contiguous(idxView); // slices are strided views; raw readback needs contiguous
-        topVals = ops.takeAlongAxis(lp, topIdx, -1);
-        for (const a of [negLp, part, idxView]) a.dispose();
-      }
-      extras = { sel, topIdx, topVals };
-    }
-    lp.dispose();
-    return { tok, extras };
+    const result = stepSampler.sample(logits3d, step);
+    return { tok: result.token, extras: result.extras };
   };
 
-  const pushHistory = (tok: MlxArray) => {
-    if (!needsTokenHistory) return;
-    if (!history) {
-      history = ops.reshape(tok, [1]);
-    } else {
-      const prev = history;
-      history = ops.concatAxis([prev, tok], 0);
-      prev.dispose();
-    }
-  };
+  const pushHistory = (tok: MlxArray): void => stepSampler.commitDevice(tok);
 
   // Decode-loop state lives at function scope so the finally can still
   // report stats and dispose in-flight arrays when the consumer
@@ -665,127 +836,301 @@ async function* generateInner(
   let prefillMs = 0;
   let tDecode = 0;
   let decodeMs = 0;
-  let generated = 0;
+  let generated = options.initialGeneratedTokens ?? 0;
+  let finishReason: GenerateStats["finishReason"];
   const forwarded: number[] = [];
   let pending: MlxArray | null = null;
   let nextPending: MlxArray | null = null;
   let pendingExtras: StepExtras | null = null;
   let nextExtras: StepExtras | null = null;
+  let decoder: MlxDecodeStep | null = null;
+  let appender: MlxTokenAppend | null = null;
   let finished = false;
   let threw = false;
+  let executionError: unknown;
+  let closeTokenZero: (() => void) | undefined;
   const makeStats = (): GenerateStats => ({
-    promptTokens: promptTokens.length,
-    cachedTokens,
+    promptTokens: options.originalPromptTokens ?? promptTokens.length,
+    cachedTokens: Math.min(cachedTokens, options.originalPromptTokens ?? promptTokens.length),
     generatedTokens: generated,
+    ...(finishReason ? { finishReason } : {}),
     prefillMs,
     decodeMs,
-    prefillTps: ((promptTokens.length - cachedTokens) / prefillMs) * 1000,
+    prefillTps: prefillMs > 0
+      ? ((promptTokens.length - cachedTokens) / prefillMs) * 1000
+      : 0,
     decodeTps: (generated / decodeMs) * 1000,
     cacheTokens: [...promptTokens, ...forwarded],
+    ...(fillOn ? { fill: options.fill!.stats } : {}),
   });
 
+  /** The one invariant that makes fill safe: the caches hold exactly
+   *  prompt + everything `forwarded` records, so a fill append must forward
+   *  ONLY the injected ids — the normal step already consumed `cur` and wrote
+   *  its KV. Forwarding [token, ...ids] would duplicate a position and
+   *  silently corrupt both the KV and PromptCache.put's key. Checked under
+   *  MLX_BUN_FILL_TRACE=1 on both sides of the append. */
+  const assertFillAlignment = (where: string): void => {
+    const offset = cache[0]!.offset;
+    const want = promptTokens.length + forwarded.length;
+    if (offset !== want)
+      throw new Error(
+        `[fill] cache misalignment ${where}: cache offset ${offset} != ` +
+        `prompt ${promptTokens.length} + forwarded ${forwarded.length}`,
+      );
+  };
+
+  /** THE apply primitive for token fast-forwarding — ONE chunked forward
+   *  carries a proposal into the KV (and the recurrent state), and BOTH
+   *  policies ride it (src/fill/proposal.ts):
+   *
+   *   - assert: the span is determined (a template scaffold). Append and move
+   *     on. No readback, no checkpoint, no rewind. The in-flight sample for
+   *     the position after `token` is dropped UNEXAMINED — a discarded
+   *     pipeline dispatch, not a rejected draft.
+   *   - verify: the span is likely (a copy from earlier in the session).
+   *     Position 0 is checked BEFORE the forward against the in-flight sample
+   *     (free, and a mismatch costs nothing — nothing has been written).
+   *     Positions 1..m-1 are checked against the argmax already sitting in
+   *     THIS forward's logits, so verification adds no pass over the weights.
+   *     The rejected tail is rewound through the same cache contract the spec
+   *     lane's rounds use, and decode resumes at the first disagreement.
+   *
+   *  Returns the ids the engine must now emit (empty = no fill happened).
+   *  Mutates nextPending/nextExtras/forwarded, exactly like the normal step. */
+  /** One trace record: the proposal next to the model's own token at every
+   *  span position (`actual[0]` = in-flight sample, `actual[j]` = argmax after
+   *  ids[j-1]). Appended as JSONL to MLX_BUN_FILL_TRACE=<file>. */
+  const traceProposal = (
+    proposal: Proposal, generatedNow: number, accepted: number, actual: number[],
+  ): void => {
+    const ids = proposal.ids;
+    let firstMismatch = -1;
+    for (let j = 0; j < ids.length && j < actual.length; j++)
+      if (actual[j] !== ids[j]) { firstMismatch = j; break; }
+    const dec = options.fill!.decode;
+    const rec: FillTraceRecord = {
+      ts: new Date().toISOString(), origin: proposal.origin, policy: proposal.policy,
+      generated: generatedNow, proposedLen: ids.length, accepted, firstMismatch,
+      proposed: ids, actual,
+      ...(dec ? { proposedText: dec(ids), actualText: dec(actual) } : {}),
+    };
+    try { appendFileSync(fillTraceFile!, JSON.stringify(rec) + "\n"); } catch { /* trace is best-effort */ }
+  };
+
+  const applyProposal = async (
+    proposal: Proposal, generatedNow: number,
+  ): Promise<number[]> => {
+    const fill = options.fill!;
+    const ids = proposal.ids;
+    const verify = proposal.policy === "verify";
+    if (verify && !verifyCapable) {
+      fill.noteVerifyUnsupported();
+      fill.commit(proposal, 0);
+      return [];
+    }
+    // The in-flight sample IS the model's own choice for position 0 (read for
+    // the trace under both policies; the served assert path never checks it).
+    const inFlight = verify || fillTraceFile ? ops.itemUint32(nextPending!) : -1;
+    if (verify) {
+      if (inFlight !== ids[0]) {
+        if (fillTraceFile) traceProposal(proposal, generatedNow, 0, [inFlight]);
+        fill.commit(proposal, 0);
+        return []; // nextPending survives; this becomes an ordinary step
+      }
+    } else {
+      fill.noteWastedSample();
+    }
+    if (fillTrace) assertFillAlignment("at fill entry");
+    nextPending!.dispose();
+    nextPending = null;
+    disposeStepExtras(nextExtras);
+    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);
+
+    let checkpointMs = 0;
+    let roundOpen = false;
+    if (verify) {
+      const t0 = performance.now();
+      for (const c of cache) rewindable(c).specRoundBegin?.();
+      roundOpen = true;
+      checkpointMs += performance.now() - t0;
+    }
+    let accepted = ids.length;
+    let hf: MlxArray | null = null;
+    let logitsAll: MlxArray | null = null;
+    try {
+      const chunkSize = (state: readonly Cache[]) => {
+        const maxChunk = appender?.maxChunkSize(state) ?? 1;
+        return fill.appendChunkSize > 0 ? Math.min(fill.appendChunkSize, maxChunk) : maxChunk;
+      };
+      hf = await appendFillHidden(
+        !verify && appender ? appender.forwardHidden.bind(appender) : graph.forwardHidden.bind(graph),
+        cache, ids, verify ? () => ids.length : chunkSize,
+      );
+      const [, Lf, Hf] = hf.shape as [number, number, number];
+      let pred: number[] | null = null;
+      if (verify || fillTraceFile) {
+        // Free logits: this forward already computed every span position's
+        // hidden state. argmax at j is the model's continuation after ids[j],
+        // so it verifies ids[j+1]. (Under assert this readback is trace-only.)
+        logitsAll = graph.projectLogits(hf, { type: "all" });
+        const argmax = ops.argmaxAxis(logitsAll, -1);
+        pred = argmax.toIntTokens();
+        argmax.dispose();
+      }
+      if (verify && pred) {
+        for (let j = 0; j + 1 < ids.length; j++) {
+          if (pred[j] !== ids[j + 1]) { accepted = j + 1; break; }
+        }
+      }
+      if (fillTraceFile && pred)
+        traceProposal(proposal, generatedNow, accepted, [inFlight, ...pred.slice(0, ids.length - 1)]);
+      const t1 = performance.now();
+      if (accepted < ids.length) {
+        // Trimmable caches drop the rejected tail; recurrent caches restore
+        // their pre-round snapshot and bit-exactly replay the accepted prefix
+        // (Cache.specRoundRollback — the same primitive src/spec/serve-loop.ts
+        // uses after its accept walk).
+        for (const c of cache) {
+          const rc = rewindable(c);
+          if (rc.specRoundRollback) rc.specRoundRollback(accepted);
+          else rc.trim(ids.length - accepted);
+        }
+        roundOpen = false;
+      } else if (roundOpen) {
+        for (const c of cache) rewindable(c).specRoundCommit?.();
+        roundOpen = false;
+      }
+      if (verify) checkpointMs += performance.now() - t1;
+      forwarded.push(...ids.slice(0, accepted));
+      if (fillTrace) assertFillAlignment("after fill append");
+      // The sampler's token history gets the ACCEPTED ids only, before the
+      // resume sample — so logits processors (repetition/presence/frequency)
+      // see exactly the sequence that was emitted.
+      stepSampler.commitNumbers(ids.slice(0, accepted));
+      // Resume from the LAST ACCEPTED position's logits: the bonus token on a
+      // full accept, the correction on a partial one. Same sampler and step
+      // index as an ordinary step — the correction's KV is written by the next
+      // iteration's forward, like any sampled token.
+      const willGen = generatedNow + accepted;
+      if (willGen < maxTokens) {
+        let logits: MlxArray;
+        if (logitsAll) {
+          const V = logitsAll.shape[2]!;
+          logits = logitsAll.slice([0, accepted - 1, 0], [1, accepted, V]);
+        } else {
+          const hLast = hf.slice([0, Lf - 1, 0], [1, Lf, Hf]);
+          logits = graph.projectLogits(hLast, { type: "all" });
+          hLast.dispose();
+        }
+        const sn = sampleStep(logits, willGen);
+        nextPending = sn.tok;
+        nextExtras = sn.extras;
+        logits.dispose();
+        ops.asyncEvalAll([nextPending, ...stepExtrasArrays(nextExtras)]);
+      }
+    } finally {
+      // A throw mid-round would otherwise leave the recurrent caches armed
+      // (their next forward raises "spec round already recorded").
+      if (roundOpen) for (const c of cache) rewindable(c).specRoundCommit?.();
+      logitsAll?.dispose();
+      hf?.dispose();
+    }
+    if (verify) fill.noteVerifyEvent(checkpointMs);
+    fill.commit(proposal, accepted);
+    if (fillTrace)
+      console.error(
+        `[fill] ${proposal.origin}/${proposal.policy} ${accepted}/${ids.length} ` +
+        `after ${generatedNow} generated (events ${fill.stats.events}, ` +
+        `injected ${fill.stats.injected}, checkpoint ${checkpointMs.toFixed(2)}ms)`,
+      );
+    return ids.slice(0, accepted);
+  };
+
+  let nextCheckpoint = options.checkpointEveryTokens && options.checkpointEveryTokens > 0
+    ? (Math.floor(generated / options.checkpointEveryTokens) + 1) * options.checkpointEveryTokens
+    : Number.POSITIVE_INFINITY;
+  let tPrefill = performance.now();
+
   try {
+    if (resuming) {
+      // The checkpoint cache already covers the complete replay prefix. Its
+      // 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]);
+      ops.asyncEvalAll([pending]);
+      prefillMs = 0;
+      closePrefill?.();
+      tDecode = performance.now();
+    } else {
     // ---- prefill ----
+    const closeInitialKv = diagnostics.trace?.begin("prefill.kv_maintenance", {
+      mechanism: diagnostics.mechanism ?? "serial",
+      boundary: "initial",
+    });
     maybeQuantizeKv(cache, options);
-    const tPrefill = performance.now();
+    closeInitialKv?.();
+    tPrefill = performance.now();
     let h0: MlxArray;
     if (options.promptEmbeddings) {
       if (cachedTokens !== 0)
         throw new Error("promptEmbeddings cannot be combined with a pre-warmed cache");
       if (needsTokenHistory)
-        history = ops.fromInt32(promptTokens, [promptTokens.length]);
+        stepSampler.seedHistory(promptTokens);
       // e2b/e4b need the spliced token ids to build per-layer inputs
       // (multimodal soft-token positions zeroed inside forwardEmbeddings).
       const embedIds = ops.fromInt32(promptTokens, [1, promptTokens.length]);
-      h0 = model.forwardEmbeddings(
+      h0 = binding.forwardEmbeddings!(
         options.promptEmbeddings, cache, options.imageMask ?? null, embedIds,
         options.multimodalMask ?? null,
       );
       embedIds.dispose();
     } else {
       let pos = cachedTokens;
-      // 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
+      const snapshotAt = 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 = await forwardHiddenForGeneration(model, 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 = await forwardHiddenForGeneration(model, ids, cache);
-        ids.dispose();
-        h.dispose(); // logits never computed for non-final chunks
-        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
-        // and the first decode step pays a one-shot reclaim stall that
-        // scales with prompt length (~800 ms after an 8k prefill —
-        // measured, scripts/decode-split.ts; the context-scaling decode
-        // gap's main term).
-        clearCache();
-        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]);
+        : undefined;
+      const tailSplit = runtime.flag("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+      const forward = graph.forwardHidden.bind(graph);
+      const maintain = () => maybeQuantizeKv(cache, options);
+      while (true) {
+        options.signal?.throwIfAborted();
+        const step = nextPrefillStep({ length: promptTokens.length, position: pos,
+          chunkSize: prefillChunkSize, tailSplit, snapshotAt });
+        if (step.kind === "final" && needsTokenHistory) stepSampler.seedHistory(promptTokens);
+        const closeChunk = diagnostics.trace?.begin("prefill.chunk", {
+          mechanism: diagnostics.mechanism ?? "serial", startToken: step.start,
+          tokens: step.end - step.start, ...(step.kind === "final" ? { final: true } : {}),
+        });
+        let hidden: MlxArray | null;
+        try { hidden = await executeMlxPrefillStep(forward, cache, promptTokens, step, maintain); }
+        finally { closeChunk?.(); }
+        pos = step.end;
+        if (hidden) { h0 = hidden; break; }
+        if (runtime.value("MLX_BUN_PREFILL_MEM_LOG") === "1")
+          console.error(`[prefill-mem] ${pos} active ${(activeMemory() / 2 ** 30).toFixed(2)} peak ${(peakMemory() / 2 ** 30).toFixed(2)}`);
+        await new Promise<void>((resolve) => setImmediate(resolve));
+        if (step.snapshot) options.onPrefillDone?.();
       }
-      // 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 = await forwardHiddenForGeneration(model, ids0, cache);
-      ids0.dispose();
     }
     if (options.snapshotAt === undefined || options.snapshotAt >= promptTokens.length)
       options.onPrefillDone?.();
+    // The trace's prefill span is the uncached model-forward portion only.
+    // Token-0 lm-head, sampling, and readback are a separate additive stage.
+    if (diagnostics.trace && runtime.value("MLX_BUN_P2R_SYNC") === "1")
+      synchronize(gpuStream);
+    closePrefill?.();
+    closeTokenZero = diagnostics.trace?.begin("token_zero.total", {
+      mechanism: diagnostics.mechanism ?? "serial",
+    });
     const [, L0, H] = h0.shape as [number, number, number];
     const hLast = h0.slice([0, L0 - 1, 0], [1, L0, H]);
     h0.dispose();
-    const logits0 = model.logitsFromHidden(hLast);
+    const logits0 = graph.projectLogits(hLast, { type: "all" });
     hLast.dispose();
     const s0 = sampleStep(logits0, 0); // token array [1] (+ optional extras)
     pending = s0.tok;
@@ -799,31 +1144,17 @@ async function* generateInner(
     // to prompt_time, not decode — replicated so cross-stack decode
     // tok/s measure the same quantity. The boundary cost is real and
     // scales with prompt length; it belongs to "having prefilled".
-    ops.asyncEvalAll([pending, ...extrasArrays(pendingExtras)]);
+    ops.asyncEvalAll([pending, ...stepExtrasArrays(pendingExtras)]);
+    }
 
     // ---- decode (pipelined) ----
-    // Compiled decode (docs/design/optimization_plan.md Phase A): replay the per-step
-    // graph in C++ instead of rebuilding it through bun:ffi every token.
-    // Bit-exact with the uncompiled path (tests/compiled-decode.test.ts);
-    // MLX_BUN_COMPILED_DECODE=0 is the kill switch / A-B lever. LoRA
-    // generations stay uncompiled (adapter weights would bake into the
-    // trace as constants). Any unsupported cache state falls back for
-    // the rest of the generation.
-    // MoE models stay uncompiled: GatherQMM lacks output_shapes in mlx
-    // 0.6.0, and shapeless replay re-infers the whole tape whenever the
-    // growing attention windows change shape (= every step). Remove this
-    // when upstream implements GatherQMM::output_shapes.
-    let compiled =
-      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;
+    // Selection and transactional fallback belong to the graph's binding.
+    decoder = options.decodePolicy?.compiledDecode === false ? null : binding.createDecode?.({
+      hasAdapters: !!options.adapters?.length, pagedKv: !!options.pagedKv,
+    }) ?? null;
+    appender = fillOn ? binding.createAppend?.({
+      hasAdapters: !!options.adapters?.length, pagedKv: !!options.pagedKv,
+    }) ?? null : 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
@@ -836,11 +1167,38 @@ async function* generateInner(
     // (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 = shouldUseGrammarJump(options);
+    const grammarJump = options.decodePolicy?.grammarJump ?? shouldUseGrammarJump(options, runtime);
+    // Yield token zero before its own decode forward. The captured runtime
+    // keeps this scheduling choice stable across consumer awaits.
+    const earlyFirstToken = runtime.value("MLX_BUN_EARLY_FIRST_TOKEN") === "1" &&
+      maxTokens > 1 && !resuming && !fillOn && !options.grammar;
     while (!stop) {
+      options.signal?.throwIfAborted();
       const cur = pending!;
-      const curExtras = pendingExtras;
-      pendingExtras = null;
+      const stepIndex = generated;
+      let curExtras = pendingExtras;
+      let yieldedEarly = false;
+      if (earlyFirstToken && stepIndex === 0) {
+        const firstToken = ops.itemUint32(cur);
+        options.signal?.throwIfAborted();
+        if (!eosTokenIds.includes(firstToken)) {
+          prefillMs = performance.now() - tPrefill;
+          closeTokenZero?.();
+          closeTokenZero = undefined;
+          tDecode = performance.now();
+          const logprobs = readExtras(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;
+          yield { token: firstToken, index: stepIndex, ...(logprobs ? { logprobs } : {}) };
+          options.signal?.throwIfAborted();
+        }
+      }
+      // Keep current extras owned by the run until readback succeeds. A
+      // forward/grammar failure must leave them reachable by finally.
       // 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,
@@ -861,11 +1219,11 @@ async function* generateInner(
       let jumpEmit: number[] | null = null;
       if (options.grammar) {
         grammarTok = ops.itemUint32(cur);
-        if (generated + 1 < maxTokens) {
+        if (stepIndex + 1 < maxTokens) {
           options.grammar.accept(grammarTok);
           await options.grammar.ready();
           if (grammarJump && !options.grammar.isTerminated) {
-            jumpEmit = options.grammar.jumpForward(maxTokens - (generated + 1));
+            jumpEmit = options.grammar.jumpForward(maxTokens - (stepIndex + 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();
@@ -886,98 +1244,113 @@ async function* generateInner(
         // 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();
-        }
+        stepSampler.commitNumbers(jumpEmit);
         const chunk = [grammarTok, ...jumpEmit];
         const ids = ops.fromInt32(chunk, [1, chunk.length]);
-        const h = await forwardHiddenForGeneration(model, ids, cache);
+        const h = await graph.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;
+        const willGen = stepIndex + 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);
+          const logits = graph.projectLogits(hLast, { type: "all" });
           hLast.dispose();
           const sn = sampleStep(logits, willGen);
           nextPending = sn.tok;
           nextExtras = sn.extras;
           logits.dispose();
-          ops.asyncEvalAll([nextPending, ...extrasArrays(nextExtras)]);
+          ops.asyncEvalAll([nextPending, ...stepExtrasArrays(nextExtras)]);
         } else {
           h.dispose(); // burst ends the generation (max_tokens or grammar done)
         }
-      } else if (generated + 1 < maxTokens && !options.grammar?.isTerminated) {
+      } else if (stepIndex + 1 < maxTokens && !options.grammar?.isTerminated) {
         maybeQuantizeKv(cache, options);
         pushHistory(cur);
         let logits: MlxArray | null = null;
         let evalWith: MlxArray[] = [];
-        if (compiled && CompiledDecode.supports(cache)) {
-          try {
-            const r = compiled.step(cur, cache);
-            logits = r.logits;
-            evalWith = r.evalWith;
-          } catch (e) {
-            // Safe to re-forward the SAME token below: a failed step is
-            // transactional — segmented mode stages ring adopts and rolls
-            // back committed js-layer writes on throw (see #stepSegmented),
-            // and buffer growth done by a partial prepare is benign for the
-            // uncompiled path (updateAndFetch re-checks capacity).
-            compiled = null;
-            console.warn(`compiled decode disabled for this generation: ${e}`);
-          }
+        const decoded = decoder?.tryStep(cur, cache);
+        if (decoded) {
+          logits = decoded.logits;
+          evalWith = decoded.evalWith;
         }
         if (!logits) {
           const ids = ops.reshape(cur, [1, 1]);
-          const h = await forwardHiddenForGeneration(model, ids, cache);
+          const h = await graph.forwardHidden(ids, cache);
           ids.dispose();
-          logits = model.logitsFromHidden(h);
+          logits = graph.projectLogits(h, { type: "all" });
           h.dispose();
         }
-        const sn = sampleStep(logits, generated + 1);
+        const sn = sampleStep(logits, stepIndex + 1);
         nextPending = sn.tok;
         nextExtras = sn.extras;
         logits.dispose();
-        ops.asyncEvalAll([nextPending, ...extrasArrays(nextExtras), ...evalWith]);
+        ops.asyncEvalAll([nextPending, ...stepExtrasArrays(nextExtras), ...evalWith]);
       }
 
       // sync-read step n's token while n+1 computes
       // 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);
+      options.signal?.throwIfAborted();
       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
         // "free" on the decode clock there too)
         prefillMs = performance.now() - tPrefill;
+        closeTokenZero?.();
+        closeTokenZero = undefined;
         tDecode = performance.now();
       }
       cur.dispose();
       pending = null;
-      generated++;
+      if (!yieldedEarly) generated++;
       // if a next-step graph was built, this token's KV entered the cache
       // (jump iterations pushed the whole chunk already)
       if (nextPending !== null && !jumpEmit) forwarded.push(token);
 
       if (eosTokenIds.includes(token)) {
-        disposeExtras(curExtras);
+        finishReason = "stop";
+        disposeStepExtras(curExtras);
+        pendingExtras = null;
         nextPending?.dispose();
         nextPending = null;
-        disposeExtras(nextExtras);
+        disposeStepExtras(nextExtras);
         nextExtras = null;
         stop = true;
       } else {
+        // ---- token fast-forwarding (K3): APPEND, then yield ----------------
+        // The engine already knows the next m tokens, so it writes them into
+        // the KV itself with an append and resumes sampling after them. No
+        // draft, no verify, no rollback: an injected token is indistinguishable
+        // to the model from one it sampled. The append happens BEFORE any of
+        // the burst's yields, which is what makes a consumer break mid-burst
+        // safe — `forwarded` already describes the cache exactly.
+        let fillEmit: number[] | null = null;
+        if (fillOn) {
+          if (jumpEmit)
+            throw new Error("fill and grammar jump-forward cannot share an iteration");
+          // A fill needs an in-flight step to ride; with nextPending null this
+          // token's KV was never written and the generation ends here. push()
+          // still runs so the session's history (and echo index) stay exact.
+          const proposal = options.fill!.push(
+            token, nextPending !== null ? maxTokens - generated : 0,
+          );
+          if (proposal && nextPending !== null) {
+            const emitted = await applyProposal(proposal, generated);
+            if (emitted.length) fillEmit = emitted;
+          } else if (proposal) {
+            options.fill!.commit(proposal, 0); // unreachable at budget 0; belt
+          }
+        }
         // readExtras before the yield: if the consumer breaks at this yield,
         // the extras are already read and disposed.
         const logprobs = readExtras(curExtras);
-        yield { token, index: generated - 1, ...(logprobs ? { logprobs } : {}) };
+        pendingExtras = null;
+        if (!yieldedEarly) yield { token, index: generated - 1, ...(logprobs ? { logprobs } : {}) };
+        options.signal?.throwIfAborted();
         // 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();
@@ -988,11 +1361,24 @@ async function* generateInner(
         // safe — `forwarded` already reflects the cache exactly.
         if (jumpEmit) {
           for (const jt of jumpEmit) {
+            options.signal?.throwIfAborted();
             generated++;
             yield { token: jt, index: generated - 1 };
             if ((generated - 1) % 256 === 0) clearCache();
           }
         }
+        // Fill burst: same one-yield-per-token shape as the jump burst, so
+        // CompletionSink / StopMatcher / StreamDecoder see an ordinary stream
+        // (a stop sequence inside an injected span fires exactly where it
+        // would have in an unfilled run).
+        if (fillEmit) {
+          for (const ft of fillEmit) {
+            options.signal?.throwIfAborted();
+            generated++;
+            yield { token: ft, index: generated - 1 };
+            if ((generated - 1) % 256 === 0) clearCache();
+          }
+        }
         if (nextPending === null) {
           stop = true;
         } else {
@@ -1001,6 +1387,21 @@ async function* generateInner(
           pendingExtras = nextExtras;
           nextExtras = null;
         }
+        if (
+          pending !== null &&
+          options.onDecodeCheckpoint &&
+          generated >= nextCheckpoint
+        ) {
+          const pendingToken = ops.itemUint32(pending);
+          await options.onDecodeCheckpoint({
+            cacheTokens: [...promptTokens, ...forwarded],
+            caches: cache,
+            generatedTokens: generated,
+            pendingToken,
+          });
+          const every = options.checkpointEveryTokens!;
+          while (nextCheckpoint <= generated) nextCheckpoint += every;
+        }
       }
     }
     decodeMs = performance.now() - tDecode;
@@ -1008,21 +1409,28 @@ async function* generateInner(
     return makeStats();
   } catch (e) {
     threw = true;
+    executionError = e;
     throw e;
   } finally {
     if (!finished) {
       pending?.dispose();
       nextPending?.dispose();
-      disposeExtras(pendingExtras);
-      disposeExtras(nextExtras);
+      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 (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
       // stats — `forwarded` only lists tokens whose KV actually entered
@@ -1032,8 +1440,3 @@ async function* generateInner(
     }
   }
 }
-
-function disposing(old: MlxArray, next: MlxArray): MlxArray {
-  old.dispose();
-  return next;
-}