mlx-bun 0.0.9 → 0.0.11

package/src/serve/batch-scheduler.ts

@@ -21,9 +21,12 @@
 // mask). Full layers share one leftPad/offset (all rows advance together); the
 // rotating caches self-track per-row leftPad/offset as the ring wraps. The
 // per-row absolute position stays consistent across both (full: offset-leftPad;
-// rot: offsetArr) — see docs/design/parallel-slots.md. Models whose caches
-// don't support the dynamic-B ops (e.g. Qwen3.5's SSMCache) never reach this
-// scheduler — GenerationGateway.willBatch gates on cache capability.
+// rot: offsetArr) — see docs/design/parallel-slots.md. Hybrid gated-DeltaNet
+// models (Qwen3.5) add "ssm" layers: SSMCache state is plain [B,...] with no
+// temporal axis and no padding (rows solo-prefill unpadded, decode feeds one
+// real token per row), so merge/filter are B-axis concat/take and the cache
+// passes through the step unwrapped (no mask, no per-row RoPE — full layers
+// carry positions). Gate: GenerationGateway.willBatch on cache capability.
 //
 // Engine mechanics (the serial decode loop's hygiene, transplanted —
 // batching-v2-plan step 3):
@@ -53,16 +56,56 @@
 
 import { MlxArray } from "../mlx/array";
 import * as ops from "../mlx/ops";
-import { clearCache } from "../mlx/ffi";
-import { KVCache, RotatingKVCache, type Cache } from "../model/gemma4-base";
-import { BatchedDecodeMaskCache, mergeKVRows, filterKVRows } from "../model/batched-mask";
+import { clearCache, Dtype } from "../mlx/ffi";
+import { flagOn } from "../flags";
+import { CompiledDecode } from "../model/compiled-decode";
+import type { Gemma4Model } from "../model/gemma4";
+import {
+  KVCache, QuantizedKVCache, RotatingKVCache, RotatingQuantizedKVCache, type Cache,
+} from "../model/gemma4-base";
+import { BatchedDecodeMaskCache, mergeKVRows, extendKVRows, extractKVRow, filterKVRows } from "../model/batched-mask";
+import {
+  BatchedQuantDecodeMaskCache, extendQuantRows, extractQuantRow, filterQuantRows, mergeQuantRows,
+  type QuantRow,
+} from "../model/batched-quant";
+import { BatchedRotatingQuantCache } from "../model/batched-rotating-quant";
+import type { KvQuantSpec } from "../config";
 import { BatchedRotatingCache } from "../model/batched-rotating";
+import { cloneKvCaches } from "../kv-store";
+import { SSMCache } from "../model/qwen3-delta";
 import type { RuntimeModel } from "../model/factory";
+import type { GrammarController } from "../grammar";
+import { kvBytesAt } from "../fit";
+import { toLogprobs } from "../sampler";
 
 /** Decode-pipeline kill switch (read once at load, like the serial loop's
  *  MLX_BUN_COMPILED_DECODE): 1 ⇒ read each step's tokens synchronously. */
 const NO_PIPELINE = process.env.MLX_BUN_BATCH_NO_PIPELINE === "1";
 
+/** Per-step phase timing (MLX_BUN_BATCH_STEP_TRACE=1, debug-only): where a
+ *  decode step's wall time goes — graph BUILD+dispatch (host), the pipelined
+ *  READ of the previous step's tokens (GPU wait), row EMIT (onToken/SSE), and
+ *  the GAP between consecutive steps (drive-loop + everything else). Sums
+ *  print via `stepTraceReport()` (the b1 profile experiment calls it). */
+const STEP_TRACE = process.env.MLX_BUN_BATCH_STEP_TRACE === "1";
+const STEP_T = { t0: 0, lastEnd: 0, build: 0, read: 0, emit: 0, gap: 0, n: 0 };
+export function stepTraceReport(): string {
+  const per = (x: number) => (STEP_T.n ? (x / STEP_T.n).toFixed(3) : "0");
+  const s = `steps=${STEP_T.n} build=${per(STEP_T.build)}ms read=${per(STEP_T.read)}ms emit=${per(STEP_T.emit)}ms gap=${per(STEP_T.gap)}ms (per step)`;
+  STEP_T.t0 = STEP_T.lastEnd = STEP_T.build = STEP_T.read = STEP_T.emit = STEP_T.gap = STEP_T.n = 0;
+  return s;
+}
+
+// NOTE — oMLX-style "adaptive burst decode" (several steps per event-loop
+// yield, omlx/engine_core.py _step_burst) was ported here 2026-07-02 and
+// REFUTED by measurement: on this runtime it REGRESSED cpm5 B=4 aggregate
+// 345→289 tok/s, batch-lane B=1 149→121, and TTFT ~+100 ms (first-token SSE
+// flush waits out the burst budget). Their win exists because each Python
+// step hand-off ping-pongs the GIL with asyncio/uvicorn (~1 ms/token); Bun's
+// setImmediate hop costs microseconds, so bursting here only delays socket
+// flushes. Don't re-add without new evidence — the per-yield step below is
+// the measured optimum (docs/design/batching-perf-path.md P4 notes).
+
 /** A token sampler for one row: (logits [1,V], step) → token array [1] on
  *  device. Greedy is `(l) => ops.argmaxAxis(l, -1)`; richer closures fold in
  *  temperature / top-p / logits processors + history (built per request from
@@ -81,17 +124,55 @@ export interface BatchRequest {
    *  THIS row only (its submit promise rejects; siblings continue). May be
    *  async; keep it cheap — it runs inline in the step loop. */
   onToken: (token: number) => void | boolean | Promise<void | boolean>;
+  /** The row's sampler is PLAIN GREEDY (temperature 0, no curve) with no
+   *  logits processors and no grammar — set by the gateway when true. Lets
+   *  the scheduler take the vectorized sampling fast path (ONE
+   *  log-softmax+argmax over [B,V] instead of B slice/sample/concat graphs)
+   *  when EVERY live row qualifies. Numerically identical per row
+   *  (row-independent ops, same per-row shapes); the per-row closure path
+   *  is the fallback and the MLX_BUN_BATCH_VEC_SAMPLE=0 kill switch. */
+  plainGreedy?: boolean;
+  /** Grammar controller for this row (B1: per-row matchers). When set, the
+   *  scheduler drives it: `accept()` after each emitted token (fires the async
+   *  bitmask fill), `await ready()` before the row's next sample (the sample
+   *  closure applies `applyMask`). Termination (grammar satisfied) is an
+   *  additional per-row stop source. The gateway OWNS disposal (finally around
+   *  submit()); the scheduler uses but never owns. Null on the degrade path. */
+  grammar?: GrammarController;
+  /** Stable cache boundary from the server's template probe (the serial
+   *  lane's options.snapshotAt): the prompt prefix the NEXT turn's render
+   *  preserves. The scheduler snapshots a trim-free prompt-cache entry at
+   *  min(this, promptIds.length-1) during the solo prefill — the oracle
+   *  invariant (mlx-lm insert_segments) that gives wrapped-ring models a
+   *  reuse path. Absent = still capped at promptIds.length-1. */
+  snapshotAt?: number;
 }
 
 export interface BatchStats {
   promptTokens: number;
   generatedTokens: number;
-  /** Prompt tokens served from a pre-warmed cache. Always 0 in v1 (each row is
-   *  solo-prefilled from scratch; prompt-cache reuse under batching is later). */
+  /** Prompt tokens served from the prompt cache (Phase 3.2): a joiner's solo
+   *  prefill starts from the longest usable cached prefix, exactly like the
+   *  serial lane's runGeneration take(). 0 on a cold prefill. */
   cachedTokens: number;
   finishReason: "stop" | "length";
 }
 
+/** The slice of PromptCache the scheduler drives (structural — the server's
+ *  PromptCache satisfies it). take() on admission (any joiner: a restored
+ *  prefix + suffix prefill is byte-safe whether the row later merges or
+ *  not). put() paths: rows that finish NEVER-MERGED (adopted lone rows,
+ *  Phase 3.2) hand their pristine serial caches over zero-copy; rows that
+ *  finish INSIDE a multi-row batch get their KV EXTRACTED per row into
+ *  fresh serial caches first (mlx-lm server.py:872 extract_cache — the
+ *  cross-request reuse concurrent agents live on). */
+export interface RowPromptCache {
+  take(prompt: number[], ns?: string): {
+    tokens: number[]; caches: Cache[]; retain?: () => void;
+  } | null;
+  put(tokens: number[], caches: Cache[], ns?: string, retain?: () => void): void;
+}
+
 interface Row {
   req: BatchRequest;
   resolve: (s: BatchStats) => void;
@@ -100,6 +181,28 @@ interface Row {
   generated: number; // tokens emitted so far (incl. a terminating EOS)
   sampled: number; // sample() calls so far (leads `generated` by 1 in-pipeline)
   promptTokens: number;
+  /** Prompt tokens restored from the prompt cache at admission (stats). */
+  cachedTokens: number;
+  /** Generated tokens whose KV actually entered the cache (fed as a step
+   *  input) — serial generate()'s `forwarded` list. promptIds+fed is the
+   *  exact token coverage of the row's caches, the put() entry key.
+   *  PER-ROW EXACT for every row: pushes are gated on the pending
+   *  register's per-slot real/placeholder flags (#pendingReal), so a
+   *  doomed slot's placeholder is never recorded — and if one ever DID
+   *  feed a forward, `fedTainted` flips instead (see #step's read). */
+  fed: number[];
+  /** A placeholder slot value fed this row's KV (junk token in the cache):
+   *  the coverage key is no longer derivable — extraction refuses.
+   *  Structurally unreachable today (a doomed row always evicts at the
+   *  emit that follows its placeholder, filtering the slot before the
+   *  next forward; cold steps never contain doomed rows because a flush
+   *  emits every in-flight token) — this flag makes `fed` provably exact
+   *  instead of resting on that timing chain. */
+  fedTainted: boolean;
+  /** Ever merged into a multi-row batch (or poisoned by a reject): its KV
+   *  is interleaved in batched buffers — a finish put() must EXTRACT the
+   *  row (#extractRowCaches) rather than adopt the inners. */
+  merged: boolean;
 }
 
 /** A joiner mid-prefill: `pos` prompt tokens are already in `solo`. Advanced
@@ -108,6 +211,13 @@ interface PrefillState {
   row: Row;
   solo: Cache[];
   pos: number;
+  /** Release hook carried by an SSD-restored prompt-cache entry (must run
+   *  only after `solo`/the adopted caches are disposed — see PromptCacheEntry). */
+  retain?: () => void;
+  /** Pending boundary-snapshot position: #prefillChunk splits its chunking
+   *  exactly here, clones + put()s the trim-free prefix entry, then nulls
+   *  it. null = no snapshot for this row (short prompt / cached past it). */
+  snapAt: number | null;
 }
 
 /** Held while the batch is active; the serial fallback acquires the same lock. */
@@ -125,9 +235,41 @@ export interface BatchSchedulerOptions {
   admissionHeld?: () => boolean;
   /** Joiner prefill chunk length (default 2048, the serial loop's constant). */
   prefillChunkSize?: number;
+  /** Aggregate KV-byte budget across all running rows (batching-perf-path
+   *  P3 admission, via `--kv-budget`). A joiner whose PROJECTED KV
+   *  (kvBytesAt(config, prompt + maxTokens), sliding window already capped)
+   *  would push the batch's projected total over this ceiling WAITS in the
+   *  queue (FIFO, head-of-line — no starvation, no reorder) until rows
+   *  evict. A request that can NEVER fit (over budget alone, empty batch)
+   *  is rejected instead of deadlocking. Unset = unlimited (v1 behavior). */
+  kvBudgetBytes?: number;
+  /** Per-layer mixed-precision KV scheme (kv_config.json entries) — Phase
+   *  3.1 batched quantized KV. Server-wide constant (one scheme per server
+   *  lifetime), so it lives here at construction, not on BatchRequest.
+   *  P1 scope: only layers whose cache is a plain full-attention KVCache
+   *  convert; the GATEWAY admits batching only when every configured
+   *  layerIdx is such a layer (rotating/uniform stay serial). Each joiner's
+   *  solo prefill converts at the SAME chunk boundaries as the serial
+   *  maybeQuantizeKv, so a row's quantized bytes are bit-exact vs serial
+   *  `--kv-quant config` by construction (the L2-oracle composition rule). */
+  kvConfig?: KvQuantSpec[];
+  /** Prompt-cache hook (Phase 3.2): admission take()s the longest usable
+   *  prefix into the joiner's solo caches (suffix-only prefill — the
+   *  multi-turn chat TTFT path); rows that finish never-merged put() their
+   *  caches back. Adapter requests never reach the batch lane, so the
+   *  namespace is always "" here. Runs under the gateway mutex domain, so
+   *  take/put never race the serial lane's use of the same cache. */
+  promptCache?: RowPromptCache;
 }
 
-type LayerInner = KVCache | BatchedRotatingCache;
+type LayerInner =
+  | KVCache
+  | QuantizedKVCache
+  | BatchedRotatingCache
+  | BatchedRotatingQuantCache // rot layer under a kv_config scheme (milestone 2)
+  | RotatingKVCache // adopted lone-row state only (see #mergeJoiner adopt)
+  | RotatingQuantizedKVCache // adopted lone-row state, quantized rot layer
+  | SSMCache;
 type Row1 = { keys: MlxArray; values: MlxArray };
 
 export class BatchScheduler {
@@ -139,37 +281,113 @@ export class BatchScheduler {
   /** Sampled-but-unread token array [B], aligned with #running — the decode
    *  pipeline register. Filtered/disposed alongside the batched KV. */
   #pendingToks: MlxArray | null = null;
+  /** Per-slot flags for #pendingToks: true = a REAL sampled token of that
+   *  row's stream; false = a doomed/terminated row's placeholder (fed-token
+   *  accounting must skip it — see Row.fed/fedTainted). Filtered/cleared in
+   *  lockstep with the register. */
+  #pendingReal: boolean[] | null = null;
   #steps = 0; // decode-step counter (clearCache cadence)
+  #lastHop = 0; // last macrotask yield (B=1 fast-path rate limiting)
   #looping = false;
   #wake: (() => void) | null = null;
   readonly #maxBatch: number;
   readonly #lock: ExclusiveLock | undefined;
   readonly #admissionHeld: (() => boolean) | undefined;
   readonly #prefillChunkSize: number;
-  readonly #kinds: ("full" | "rot")[]; // per-layer attention type
+  readonly #kvBudgetBytes: number | undefined;
+  readonly #promptCache: RowPromptCache | undefined;
+  /** Retain hook of the currently-ADOPTED row's cache entry (at most one:
+   *  only a lone adopted row holds un-copied entry caches). Runs after the
+   *  adopted caches are disposed, or transfers back on put(). */
+  #adoptedRetain: (() => void) | null = null;
+  readonly #kinds: ("full" | "rot" | "ssm")[]; // per-layer attention type
   readonly #rotMaxSize: number[]; // per-layer sliding window (rot layers only)
+  /** layerIdx → mixed-precision spec (Phase 3.1); null = bf16 batch (v1). */
+  readonly #kvByLayer: Map<number, KvQuantSpec> | null;
+  /** Compiled decode runner for the B=1 serial-class case (Phase 3.2) —
+   *  same eligibility gate as generate.ts (gemma dense, kill switch
+   *  MLX_BUN_COMPILED_DECODE); adapters never reach the batch lane. Set
+   *  to null permanently on a failed step (serial disables per
+   *  generation; the scheduler is one long-lived "generation"). */
+  #compiled: CompiledDecode | null;
 
   constructor(private readonly model: RuntimeModel, opts: BatchSchedulerOptions) {
     this.#maxBatch = Math.max(1, Math.floor(opts.maxBatch));
     this.#lock = opts.lock;
     this.#admissionHeld = opts.admissionHeld;
     this.#prefillChunkSize = Math.max(1, Math.floor(opts.prefillChunkSize ?? 2048));
+    this.#kvBudgetBytes = opts.kvBudgetBytes;
+    this.#promptCache = opts.promptCache;
     const proto = model.makeCache(); // fresh caches hold no buffers
-    this.#kinds = proto.map((c) => (c instanceof RotatingKVCache ? "rot" : "full"));
+    this.#kinds = proto.map((c) =>
+      c instanceof RotatingKVCache ? "rot" : c instanceof SSMCache ? "ssm" : "full",
+    );
+    this.#kvByLayer = opts.kvConfig?.length
+      ? new Map(opts.kvConfig.map((e) => [e.layerIdx, e]))
+      : null;
     this.#rotMaxSize = proto.map((c) => (c instanceof RotatingKVCache ? c.maxSize : 0));
     for (const c of proto) c.dispose();
+    this.#compiled =
+      flagOn("MLX_BUN_COMPILED_DECODE", true) &&
+      model.config.modelType.startsWith("gemma4") &&
+      !model.config.text.enableMoeBlock
+        ? CompiledDecode.for(model as Gemma4Model)
+        : null;
   }
 
   get activeRows(): number {
     return this.#running.length;
   }
 
+  get pendingRows(): number {
+    return this.#pending.length + (this.#prefill ? 1 : 0);
+  }
+
+  /** Projected KV bytes of one row at its worst case (full prompt + full
+   *  completion; the sliding-window term is window-capped by kvBytesAt). */
+  #rowKvBytes(row: Row): number {
+    return kvBytesAt(this.model.config, row.promptTokens + row.req.maxTokens);
+  }
+
+  /** Projected aggregate KV of everything admitted (running + mid-prefill). */
+  get projectedKvBytes(): number {
+    let total = this.#running.reduce((a, r) => a + this.#rowKvBytes(r), 0);
+    if (this.#prefill) total += this.#rowKvBytes(this.#prefill.row);
+    return total;
+  }
+
+  get kvBudgetBytes(): number | undefined {
+    return this.#kvBudgetBytes;
+  }
+
+  /** KV-budget admission for the queue head. True = admit now. A candidate
+   *  that cannot fit even alone is rejected here (never deadlocks the
+   *  queue); one that fits alone but not alongside the current batch waits. */
+  #kvAdmits(candidate: Row): boolean {
+    if (this.#kvBudgetBytes === undefined) return true;
+    const need = this.#rowKvBytes(candidate);
+    if (need > this.#kvBudgetBytes && this.#running.length === 0 && !this.#prefill) {
+      this.#pending.shift();
+      candidate.reject(
+        new Error(
+          `kv budget: request needs ~${(need / 1e9).toFixed(2)} GB KV ` +
+            `(prompt ${candidate.promptTokens} + max_tokens ${candidate.req.maxTokens}), ` +
+            `over --kv-budget ${(this.#kvBudgetBytes / 1e9).toFixed(2)} GB — ` +
+            `lower max_tokens or raise the budget`,
+        ),
+      );
+      return false;
+    }
+    return this.projectedKvBytes + need <= this.#kvBudgetBytes;
+  }
+
   /** Submit a request; resolves when its row finishes (EOS, stop, or length). */
   submit(req: BatchRequest): Promise<BatchStats> {
     return new Promise<BatchStats>((resolve, reject) => {
       this.#pending.push({
         req, resolve, reject,
         current: 0, generated: 0, sampled: 0, promptTokens: req.promptIds.length,
+        cachedTokens: 0, fed: [], fedTainted: false, merged: false,
       });
       this.#ensureLoop();
     });
@@ -209,8 +427,31 @@ export class BatchScheduler {
         // Start at most one joiner prefill (drain-aware); it advances one
         // chunk per iteration, interleaved with the decode step below.
         if (!this.#prefill && !held &&
-            this.#pending.length > 0 && this.#running.length < this.#maxBatch)
-          this.#prefill = { row: this.#pending.shift()!, solo: this.model.makeCache(), pos: 0 };
+            this.#pending.length > 0 && this.#running.length < this.#maxBatch &&
+            this.#kvAdmits(this.#pending[0]!)) {
+          const row = this.#pending.shift()!;
+          // Prompt-cache reuse (Phase 3.2): start the solo prefill from the
+          // longest usable cached prefix — the serial lane's take(), same
+          // semantics (entry ownership transfers; put back or dispose).
+          const hit = this.#promptCache?.take(row.req.promptIds) ?? null;
+          if (hit) row.cachedTokens = hit.tokens.length;
+          // Boundary snapshot (the serial lane's A1, ported): a trim-free
+          // strict-prefix entry at min(stableLen, len-1) for every
+          // substantial cold prefill — same gates as runGeneration
+          // (>= 256 tokens, extends past the cached prefix). Without this
+          // the batch lane's ONLY entry was the untrimmable [prompt+gen]
+          // finish-time put — a total ctx-repeat miss once rings wrap
+          // (12B batched 84 s vs serial 0.4 s, 2026-07-06).
+          const len = row.req.promptIds.length;
+          const boundary = Math.min(row.req.snapshotAt ?? len, len - 1);
+          const snapAt =
+            this.#promptCache && boundary >= 256 && boundary > (hit?.tokens.length ?? 0)
+              ? boundary
+              : null;
+          this.#prefill = hit
+            ? { row, solo: hit.caches, pos: hit.tokens.length, retain: hit.retain, snapAt }
+            : { row, solo: this.model.makeCache(), pos: 0, snapAt };
+        }
 
         if (this.#prefill) {
           const p = this.#prefill;
@@ -221,6 +462,7 @@ export class BatchScheduler {
             // the running batch.
             this.#prefill = null;
             for (const c of p.solo) c.dispose();
+            p.retain?.();
             p.row.reject(e);
           }
         }
@@ -231,7 +473,9 @@ export class BatchScheduler {
         // still interleave: a mid-prefill joiner leaves #prefill non-null, so
         // the loop falls through to run a decode step between chunks.
         if (this.#prefill === null && !held &&
-            this.#pending.length > 0 && this.#running.length < this.#maxBatch)
+            this.#pending.length > 0 && this.#running.length < this.#maxBatch &&
+            (this.#kvBudgetBytes === undefined ||
+              this.projectedKvBytes + this.#rowKvBytes(this.#pending[0]!) <= this.#kvBudgetBytes))
           continue;
 
         if (this.#running.length > 0) {
@@ -240,10 +484,25 @@ export class BatchScheduler {
           } catch (e) {
             // A forward/sampling error can't be attributed to one row — drop
             // the batch. (A row's onToken error is contained in #emitRows.)
+            // dropOnly: mid-step caches may be inconsistent — never put().
             for (const row of this.#running) row.reject(e);
-            this.#applyFilter([]);
+            this.#applyFilter([], true);
           }
         }
+        // The per-step macrotask hop exists so joiners/admin interleave at
+        // token granularity. With ONE row and an empty queue there is
+        // nobody to interleave — rate-limit the hop to the serial loop's
+        // 25 ms cadence (identical responsiveness budget; generate()'s
+        // hoppingOnToken does exactly this). This was the B=1 drive-loop
+        // tax: one setImmediate round-trip per token ≈ the 0.4–0.8% gap
+        // vs pure serial (paired-measured 2026-07-05).
+        if (
+          this.#running.length === 1 && this.#pending.length === 0 &&
+          !this.#prefill && this.#admissionHeld?.() !== true &&
+          performance.now() - this.#lastHop < 25
+        )
+          continue;
+        this.#lastHop = performance.now();
         await new Promise<void>((r) => setImmediate(r));
       }
     } finally {
@@ -252,28 +511,111 @@ export class BatchScheduler {
     }
   }
 
+  /** Per-layer mixed-precision conversion of a joiner's SOLO caches — the
+   *  scheduler-side mirror of the serial maybeQuantizeKv (generate.ts): same
+   *  per-layer map, same skip rules (empty cache, already quantized), same
+   *  streaming discipline (evalAll the converted layer's state so the bf16
+   *  source frees before the next layer converts). Called at every prefill
+   *  chunk boundary AND once before merge, exactly where the serial loop
+   *  calls maybeQuantizeKv — that placement is what makes a row's quantized
+   *  bytes bit-exact vs serial `--kv-quant config`. P1: plain KVCache
+   *  layers only (the gateway guarantees the config maps only to those). */
+  #quantizeSolo(solo: Cache[]): void {
+    if (!this.#kvByLayer) return;
+    for (let i = 0; i < solo.length; i++) {
+      const e = this.#kvByLayer.get(i);
+      const c = solo[i]!;
+      if (!e || c.offset === 0) continue;
+      // Milestone 2: rotating layers convert too — RotatingKVCache.
+      // toQuantized → RotatingQuantizedKVCache, exactly maybeQuantizeKv's
+      // dispatch (generate.ts). Converted caches don't match either class
+      // (the four cache classes are siblings), so re-conversion never fires.
+      if (!(c instanceof KVCache || c instanceof RotatingKVCache)) continue;
+      const q = c.toQuantized(e.groupSize, e.bits);
+      solo[i] = q;
+      ops.evalAll(q.state());
+      clearCache();
+    }
+  }
+
   /** Advance a joiner's solo prefill by one chunk. Non-final chunks forward +
    *  eval the cache and return false (the caller interleaves a decode step).
    *  The final chunk samples token 0, emits it, and — if the row survives —
    *  merges it into the running batch; returns true (admission complete). */
   async #prefillChunk(p: PrefillState): Promise<boolean> {
     const prompt = p.row.req.promptIds;
-    if (prompt.length - p.pos > this.#prefillChunkSize) {
+    // Boundary-snapshot chunking (mirrors generate.ts's snapshotAt split):
+    // chunk edges land EXACTLY on snapAt so the clone is taken while the
+    // caches hold precisely that prefix — slicing tokens against caches at
+    // any other offset would store corrupt entries. The extra split only
+    // changes chunk shapes for opted-in rows: the same numerics class as a
+    // prompt-cache-hit continuation prefill (GEMV/GEMM reduction order).
+    if (p.snapAt !== null && p.pos < p.snapAt) {
+      const end = Math.min(p.pos + this.#prefillChunkSize, p.snapAt);
+      const chunk = prompt.slice(p.pos, end);
+      const ids = ops.fromInt32(chunk, [1, chunk.length]);
+      const h = this.model.forwardHidden(ids, p.solo);
+      ids.dispose();
+      h.dispose();
+      ops.evalAll(p.solo.flatMap((c) => c.state()));
+      this.#quantizeSolo(p.solo); // serial converts at every chunk boundary
+      clearCache();
+      p.pos = end;
+      if (p.pos === p.snapAt) {
+        try {
+          this.#promptCache!.put(prompt.slice(0, p.snapAt), cloneKvCaches(p.solo));
+        } catch (e) {
+          console.warn(`batch-lane boundary snapshot skipped: ${(e as Error).message}`);
+        }
+        p.snapAt = null;
+      }
+      return false;
+    }
+    // Oracle prefill convention (see generate.ts prefill): drain only to
+    // len-1 in chunks, then step-0 logits from an L=1 forward of the LAST
+    // prompt token — mirrors mlx-lm's batched engine (insert_segments'
+    // forced final 1-token segment + GenerationBatch._step,
+    // generate.py:1645/1327) and keeps this lane bit-exact with the serial
+    // lane (unified == --batch 1 == oracle). The L=1 step runs on the solo
+    // caches BEFORE merge, so per-row serial equivalence is by construction.
+    // MLX_BUN_PREFILL_TAIL_SPLIT=0 restores the old convention.
+    const tailSplit = flagOn("MLX_BUN_PREFILL_TAIL_SPLIT", true);
+    if ((tailSplit ? prompt.length - p.pos - 1 : prompt.length - p.pos) >
+        this.#prefillChunkSize) {
       const chunk = prompt.slice(p.pos, p.pos + this.#prefillChunkSize);
       const ids = ops.fromInt32(chunk, [1, chunk.length]);
       const h = this.model.forwardHidden(ids, p.solo);
       ids.dispose();
       h.dispose(); // logits never computed for non-final chunks
       ops.evalAll(p.solo.flatMap((c) => c.state()));
+      this.#quantizeSolo(p.solo); // serial converts at every chunk boundary
       clearCache(); // serial prefill's per-chunk clear (generate.ts)
       p.pos += this.#prefillChunkSize;
       return false;
     }
 
+    // Tail split: forward the remaining head [p.pos, len-1) as a plain chunk
+    // (serial chunk-boundary semantics: eval + quantize + clear), leaving
+    // exactly one token for the L=1 step-0 forward below. Head + step share
+    // this tick — same chunk-tick count as before, plus one L=1 dispatch
+    // (~a single decode step) in the admission tick's no-yield window.
+    if (tailSplit && p.pos < prompt.length - 1) {
+      const head = prompt.slice(p.pos, prompt.length - 1);
+      const ids = ops.fromInt32(head, [1, head.length]);
+      const h = this.model.forwardHidden(ids, p.solo);
+      ids.dispose();
+      h.dispose();
+      ops.evalAll(p.solo.flatMap((c) => c.state()));
+      this.#quantizeSolo(p.solo);
+      clearCache();
+      p.pos = prompt.length - 1;
+    }
+
     // Final chunk: slice the LAST hidden position BEFORE the lm_head — running
     // logitsFromHidden on the whole [1,Lp,H] hidden materializes a [1,Lp,V]
     // transient (~4.3 GB bf16 at Gemma V=262k, 8k prompt). Same reorder as the
-    // serial path (generate.ts prefill).
+    // serial path (generate.ts prefill). Under tailSplit this chunk is exactly
+    // the last prompt token (L=1).
     const chunk = prompt.slice(p.pos);
     const ids = ops.fromInt32(chunk, [1, chunk.length]);
     const h = this.model.forwardHidden(ids, p.solo);
@@ -292,12 +634,38 @@ export class BatchScheduler {
     p.row.generated = 1;
     clearCache(); // drop the prefill transients (serial's token-0 clear)
 
+    // Grammar (B1): the mask0 was applied inside the sample closure (the
+    // controller is primed at compile). After reading token 0, advance the
+    // matcher — fires the async fill for the NEXT step, overlapping the merge
+    // below. If the grammar is already satisfied at token 0 (a 1-token grammar,
+    // e.g. guided_choice landing on a single-token option), emit + finish("stop")
+    // WITHOUT merging into the batch — the row never joins #running.
+    if (p.row.req.grammar) {
+      p.row.req.grammar.accept(tok);
+      if (p.row.req.grammar.isTerminated) {
+        const stop = await this.#emit(p.row, tok);
+        // Token 0 was sampled but never fed — the caches cover exactly the
+        // prompt, a clean prompt-only entry (put-or-dispose).
+        this.#putOrDispose(p.solo, p.row.req.promptIds, p.retain);
+        this.#finish(p.row, stop === "continue" ? "stop" : stop);
+        return true;
+      }
+    }
+
     const stop = await this.#emit(p.row, tok);
     if (stop !== "continue") {
-      for (const c of p.solo) c.dispose();
+      this.#putOrDispose(p.solo, p.row.req.promptIds, p.retain);
       this.#finish(p.row, stop);
       return true;
     }
+    // Default tail-split path: the caches were already converted at the head
+    // boundary above (oracle composition: prefill ids[:-1] → convert → L=1
+    // step-0), so this call is an idempotent no-op (converted caches match
+    // neither serial class). It is LOAD-BEARING only under
+    // MLX_BUN_PREFILL_TAIL_SPLIT=0 — the old serial order: token 0 sampled
+    // from the unconverted final-chunk logits, THEN the caches convert
+    // (before decode step 1 == before the merge).
+    this.#quantizeSolo(p.solo);
     await this.#mergeJoiner(p);
     return true;
   }
@@ -308,19 +676,154 @@ export class BatchScheduler {
   async #mergeJoiner(p: PrefillState): Promise<void> {
     await this.#flushPipeline();
 
+    // ADOPT, don't copy (unified-engine plan Phase 3.2): a row joining an
+    // EMPTY batch keeps its solo caches as the batch inners — a pointer
+    // handoff, zero bytes moved (the old path ran the full merge machinery
+    // to produce a byte-identical [1,...] copy). The copy now happens only
+    // when a SECOND row joins and a genuinely new layout must exist. The
+    // prize beyond the saved copy: the lone row's caches stay SERIAL-CLASS
+    // (KVCache / RotatingKVCache / QuantizedKVCache), so the B=1 step is
+    // literally the serial graph, and compiled decode + prompt-cache
+    // take/put become possible for it. The rot branch below knows how to
+    // treat an adopted RotatingKVCache as the merge's first row.
+    if (!this.#inners) {
+      this.#inners = p.solo as LayerInner[];
+      this.#fullLeftPad = [0];
+      this.#adoptedRetain = p.retain ?? null;
+      this.#running.push(p.row);
+      return;
+    }
+
     const prev = this.#inners;
     const prevPad = this.#fullLeftPad;
     const B = this.#running.length;
     const newInners: LayerInner[] = [];
     let newFullPad = this.#fullLeftPad;
     for (let layer = 0; layer < this.#kinds.length; layer++) {
+      if (this.#kinds[layer] === "ssm") {
+        // No temporal axis, no left-pad: B-axis concat of the state slots.
+        // mergeRows steals the solo arrays when the batch starts cold, so the
+        // unconditional p.solo dispose below stays safe either way.
+        newInners.push(SSMCache.mergeRows(
+          (prev?.[layer] as SSMCache | undefined) ?? null,
+          p.solo[layer] as SSMCache,
+        ));
+        continue;
+      }
+      if (p.solo[layer] instanceof QuantizedKVCache) {
+        // Phase 3.1 — quantized full layer: same merge/extend shapes as the
+        // bf16 branch below, over (packed, scales, biases) triples. The solo
+        // row was converted by #quantizeSolo with the serial ops, so its
+        // bytes already bit-match serial `--kv-quant config`; this branch
+        // only re-arranges rows along the batch axis.
+        const qSolo = p.solo[layer] as QuantizedKVCache;
+        const [qk, qv] = qSolo.temporalView();
+        const qRow = { keys: qk, values: qv };
+        const dispose3 = (t: { packed: MlxArray; scales: MlxArray; biases: MlxArray }) => {
+          t.packed.dispose(); t.scales.dispose(); t.biases.dispose();
+        };
+        const prevQ = prev?.[layer] as QuantizedKVCache | undefined;
+        if (prevQ && process.env.MLX_BUN_BATCH_EXTEND !== "0") {
+          const [k0, v0] = prevQ.temporalView();
+          const ext = extendQuantRows(k0, v0, prevPad, qRow);
+          dispose3(k0); dispose3(v0);
+          newFullPad = ext.leftPad;
+          const c = new QuantizedKVCache(qSolo.groupSize, qSolo.bits);
+          c.restoreState(ext.keys, ext.values, ext.width);
+          newInners.push(c);
+        } else {
+          const rows: { keys: typeof qk; values: typeof qv }[] = [];
+          if (prevQ) {
+            const [k0, v0] = prevQ.temporalView(); // [B,H,off,*]
+            const S = k0.packed.shape[2]!;
+            for (let b = 0; b < B; b++) {
+              const pad = prevPad[b]!;
+              const cutRow = (t: typeof k0): typeof k0 => ({
+                packed: t.packed.slice([b, 0, pad, 0], [b + 1, t.packed.shape[1]!, S, t.packed.shape[3]!]),
+                scales: t.scales.slice([b, 0, pad, 0], [b + 1, t.scales.shape[1]!, S, t.scales.shape[3]!]),
+                biases: t.biases.slice([b, 0, pad, 0], [b + 1, t.biases.shape[1]!, S, t.biases.shape[3]!]),
+              });
+              rows.push({ keys: cutRow(k0), values: cutRow(v0) });
+            }
+            dispose3(k0); dispose3(v0);
+          }
+          rows.push(qRow);
+          const merged = mergeQuantRows(rows);
+          newFullPad = merged.leftPad;
+          const c = new QuantizedKVCache(qSolo.groupSize, qSolo.bits);
+          c.restoreState(merged.keys, merged.values, merged.width);
+          newInners.push(c);
+          for (const r of rows) { dispose3(r.keys); dispose3(r.values); }
+        }
+        // qRow views are disposed via the rows loop above or here for extend
+        if (prevQ && process.env.MLX_BUN_BATCH_EXTEND !== "0") { dispose3(qRow.keys); dispose3(qRow.values); }
+        else if (!prevQ) { /* disposed in the rows loop */ }
+        continue;
+      }
+      // NOTE: the quantized-rotating branch below never uses this layer's
+      // bf16 temporalView — calling it before the branch leaked the pair
+      // (six arrays per second-row join on quantized rotating layers,
+      // introduced in 859572d; 2026-07-07 review fix). The view is taken
+      // AFTER the branch, on the paths that actually consume it.
       const soloC = p.solo[layer] as KVCache | RotatingKVCache;
+      if (this.#kinds[layer] === "rot" && p.solo[layer] instanceof RotatingQuantizedKVCache) {
+        // Milestone 2 — QUANTIZED rotating layer: the solo row converted at
+        // the serial boundaries (#quantizeSolo), so its ring bytes are the
+        // serial oracle's; this branch re-arranges temporal triples across
+        // the batch axis (the bf16 rot merge over triples).
+        const qSolo = p.solo[layer] as RotatingQuantizedKVCache;
+        const [qk, qv] = qSolo.temporalView();
+        const rows: QuantRow[] = [];
+        const offsets: number[] = [];
+        const dispose3 = (t: QuantRow) => {
+          t.keys.packed.dispose(); t.keys.scales.dispose(); t.keys.biases.dispose();
+          t.values.packed.dispose(); t.values.scales.dispose(); t.values.biases.dispose();
+        };
+        const prevC = prev?.[layer];
+        if (prevC instanceof BatchedRotatingQuantCache) {
+          const [k0, v0] = prevC.temporalView(); // [B,H,valid,*] triples, temporal
+          const valid = k0.packed.shape[2]!;
+          for (let b = 0; b < B; b++) {
+            const pad = Math.max(0, prevC.leftPad[b]!);
+            const cutRow = (t: typeof k0): typeof k0 => ({
+              packed: t.packed.slice([b, 0, pad, 0], [b + 1, t.packed.shape[1]!, valid, t.packed.shape[3]!]),
+              scales: t.scales.slice([b, 0, pad, 0], [b + 1, t.scales.shape[1]!, valid, t.scales.shape[3]!]),
+              biases: t.biases.slice([b, 0, pad, 0], [b + 1, t.biases.shape[1]!, valid, t.biases.shape[3]!]),
+            });
+            rows.push({ keys: cutRow(k0), values: cutRow(v0) });
+            offsets.push(prevC.offsetArr[b]!);
+          }
+          for (const t of [k0, v0]) { t.packed.dispose(); t.scales.dispose(); t.biases.dispose(); }
+        } else if (prevC instanceof RotatingQuantizedKVCache) {
+          // Adopted lone row (Phase 3.2 adopt): its chronological triples
+          // are the merge's first row, pad 0 by definition.
+          const [k0, v0] = prevC.temporalView();
+          rows.push({ keys: k0, values: v0 });
+          offsets.push(prevC.offset);
+        }
+        rows.push({ keys: qk, values: qv });
+        offsets.push(qSolo.offset);
+        newInners.push(BatchedRotatingQuantCache.merge(
+          rows, offsets, this.#rotMaxSize[layer]!, qSolo.groupSize, qSolo.bits,
+        ));
+        for (const r of rows) dispose3(r);
+        continue;
+      }
       const [sk, sv] = soloC.temporalView();
       const newRow: Row1 = { keys: sk, values: sv };
       if (this.#kinds[layer] === "rot") {
         const rows: Row1[] = [];
         const offsets: number[] = [];
-        const prevRot = prev?.[layer] as BatchedRotatingCache | undefined;
+        const prevC = prev?.[layer];
+        if (prevC instanceof RotatingKVCache) {
+          // Adopted lone row (Phase 3.2): a plain serial rotating cache —
+          // its chronological view is the merge's first row, same as a
+          // fresh solo (pad 0 by definition).
+          const [k0, v0] = prevC.temporalView();
+          rows.push({ keys: k0, values: v0 });
+          offsets.push(prevC.offset);
+        }
+        const prevRot = prevC instanceof BatchedRotatingCache ? prevC : undefined;
         if (prevRot) {
           const [k0, v0] = prevRot.temporalView(); // [B,H,valid,D]
           const [, H, valid, D] = k0.shape as [number, number, number, number];
@@ -340,32 +843,58 @@ export class BatchScheduler {
         newInners.push(BatchedRotatingCache.merge(rows, offsets, this.#rotMaxSize[layer]!));
         for (const r of rows) { r.keys.dispose(); r.values.dispose(); }
       } else {
-        const rows: Row1[] = [];
         const prevFull = prev?.[layer] as KVCache | undefined;
-        if (prevFull) {
+        if (prevFull && process.env.MLX_BUN_BATCH_EXTEND !== "0") {
+          // extend-join (mlx-lm BatchKVCache.extend semantics, P0): append
+          // the new right-justified row to the running buffer in ONE pad +
+          // ONE concat — no per-row extraction. Existing pads grow, never
+          // shrink (the re-merge below re-normalizes them instead; both are
+          // masked, both token-exact vs their mlx-lm protocol twin).
+          // MLX_BUN_BATCH_EXTEND=0 = the O(B·S) re-merge, kill switch/A-B.
           const [k0, v0] = prevFull.temporalView(); // [B,H,off,D]
-          const [, H, off, D] = k0.shape as [number, number, number, number];
-          const vD = v0.shape[3]!;
-          for (let b = 0; b < B; b++) {
-            const pad = prevPad[b]!;
-            rows.push({
-              keys: k0.slice([b, 0, pad, 0], [b + 1, H, off, D]),
-              values: v0.slice([b, 0, pad, 0], [b + 1, H, off, vD]),
-            });
-          }
+          const ext = extendKVRows(k0, v0, prevPad, newRow);
           k0.dispose(); v0.dispose();
+          newFullPad = ext.leftPad;
+          const c = new KVCache();
+          c.restoreState(ext.keys, ext.values, ext.width);
+          newInners.push(c);
+          newRow.keys.dispose(); newRow.values.dispose();
+        } else {
+          const rows: Row1[] = [];
+          if (prevFull) {
+            const [k0, v0] = prevFull.temporalView(); // [B,H,off,D]
+            const [, H, off, D] = k0.shape as [number, number, number, number];
+            const vD = v0.shape[3]!;
+            for (let b = 0; b < B; b++) {
+              const pad = prevPad[b]!;
+              rows.push({
+                keys: k0.slice([b, 0, pad, 0], [b + 1, H, off, D]),
+                values: v0.slice([b, 0, pad, 0], [b + 1, H, off, vD]),
+              });
+            }
+            k0.dispose(); v0.dispose();
+          }
+          rows.push(newRow);
+          const merged = mergeKVRows(rows);
+          newFullPad = merged.leftPad;
+          const c = new KVCache();
+          c.restoreState(merged.keys, merged.values, merged.width);
+          newInners.push(c);
+          for (const r of rows) { r.keys.dispose(); r.values.dispose(); }
         }
-        rows.push(newRow);
-        const merged = mergeKVRows(rows);
-        newFullPad = merged.leftPad;
-        const c = new KVCache();
-        c.restoreState(merged.keys, merged.values, merged.width);
-        newInners.push(c);
-        for (const r of rows) { r.keys.dispose(); r.values.dispose(); }
       }
     }
     if (prev) for (const c of prev) c.dispose();
+    // An adopted row's entry-backed arrays are gone after the prev dispose;
+    // run its retain now. The joiner's likewise after its solo dispose.
+    this.#adoptedRetain?.();
+    this.#adoptedRetain = null;
     for (const c of p.solo) c.dispose();
+    p.retain?.();
+    // Every row in a REAL merge has its KV interleaved in batched buffers —
+    // no longer prompt-cache put() candidates.
+    for (const r of this.#running) r.merged = true;
+    p.row.merged = true;
     this.#inners = newInners;
     this.#fullLeftPad = newFullPad;
     this.#running.push(p.row);
@@ -381,6 +910,11 @@ export class BatchScheduler {
    *  step; filter drops the row (mlx-lm behaves identically). Length-finished
    *  rows are known in advance and are NOT sampled (placeholder slot). */
   async #step(): Promise<void> {
+    if (STEP_TRACE) {
+      const now = performance.now();
+      if (STEP_T.lastEnd) STEP_T.gap += now - STEP_T.lastEnd;
+      STEP_T.t0 = now;
+    }
     const rows = this.#running;
     const B = rows.length;
     const inners = this.#inners!;
@@ -388,55 +922,166 @@ export class BatchScheduler {
     // A row is live if it still needs tokens sampled; a row whose pending
     // unread token is its last (sampled == maxTokens) only awaits emission.
     const anyLive = rows.some((r) => r.sampled < r.req.maxTokens);
+    // Grammar (B1): if any row has a LIVE grammar controller (not terminated,
+    // still sampling), take the read-before-build shape — the matcher's
+    // acceptToken needs the token id as a JS number, so the previous step's
+    // [B] token array is read back NOW (before building the next graph),
+    // accept()ed per row (firing async fills that overlap the forward), then
+    // ready() is awaited before this step's sample. This is the serial loop's
+    // grammar resolution transplanted to the batch. Batches with NO live
+    // grammar row keep the pipelined path byte-identical (zero cost when
+    // unused). The trade: while a grammar row is live the batch runs
+    // effectively NO_PIPELINE (readback no longer overlaps GPU compute) —
+    // bounded by the readback (~0.1 ms) + fills (0.004–0.19 ms/row, overlapped
+    // with the graph build). Serial grammar pays the identical trade today.
+    const hasLiveGrammar = anyLive && rows.some(
+      (r) => r.req.grammar && !r.req.grammar.isTerminated,
+    );
+    if (hasLiveGrammar) return this.#stepGrammar();
+
     let nextToks: MlxArray | null = null;
+    let nextReal: boolean[] | null = null;
     if (anyLive) {
+      // Fed-token accounting (prompt-cache put): a COLD step feeds each
+      // row's `current` (values known now, always a real emitted token); a
+      // pipelined step feeds the pending array, whose values are pushed at
+      // the read below gated on the per-slot real flags.
+      if (!this.#pendingToks) for (const r of rows) r.fed.push(r.current);
       // Per-layer forward cache: rot layers use the persistent
-      // BatchedRotatingCache directly; full layers get a fresh
-      // BatchedDecodeMaskCache wrapper.
-      const fwd: Cache[] = inners.map((c) =>
-        c instanceof BatchedRotatingCache ? c : new BatchedDecodeMaskCache(c, B, this.#fullLeftPad, null),
-      );
+      // BatchedRotatingCache directly; ssm layers are already [B,...] state
+      // with no padding (no mask, no per-row RoPE — pass through); full
+      // layers get a fresh BatchedDecodeMaskCache wrapper — UNLESS no row
+      // has left padding. UNPADDED FAST PATH (unified-engine plan Phase 2,
+      // the B=1 case above all): with every leftPad 0 the wrapper's two
+      // jobs vanish — the padding mask (KVCache.makeMask(1) is the empty
+      // mask, exactly the serial loop's) and the per-row rope positions
+      // (every row sits at the shared scalar offset). The bare cache then
+      // dispatches the SAME per-step graph serial builds; the wrapper
+      // otherwise costs a host mask build + ~8 device nodes PER FULL LAYER
+      // PER TOKEN (the Phase-0 constant ~4–6 ms/step host tax at B=1).
+      const unpadded = this.#fullLeftPad.every((p) => p === 0);
+      // Compiled decode at B=1 (Phase 3.2): after adopt-don't-copy, a lone
+      // row's caches are SERIAL-CLASS, so the serial engine's compiled step
+      // replays the same C++ graph here — closing the batch lane's last
+      // B=1 host-tax gap (e4b's ~7%). Guards: gemma dense (constructor),
+      // serial-class caches (supports — a merged batch's BatchedRotating
+      // layers fail it), unpadded, and a uint32 pipeline register (the
+      // trace signature; per-row int32 samplers take the graph path).
+      // Grammar batches use #stepGrammar and stay on the graph path.
+      let lg: MlxArray | null = null;
+      let evalWith: MlxArray[] = [];
+      if (
+        this.#compiled && B === 1 && unpadded &&
+        (!this.#pendingToks || this.#pendingToks.dtype === Dtype.uint32) &&
+        // A filtered-to-one BATCHED rot-quant cache subclasses the serial
+        // class (so supports() passes) but carries batched ring state —
+        // exclude it; only truly-adopted serial caches replay compiled.
+        !inners.some((c) => c instanceof BatchedRotatingQuantCache) &&
+        CompiledDecode.supports(inners as Cache[])
+      ) {
+        let cur = this.#pendingToks;
+        let owned = false;
+        if (!cur) {
+          const i = ops.fromInt32([rows[0]!.current], [1]);
+          cur = i.astype(Dtype.uint32);
+          i.dispose();
+          owned = true;
+        }
+        try {
+          const r = this.#compiled.step(cur, inners as Cache[]);
+          lg = r.logits; // [1,1,V]
+          evalWith = r.evalWith;
+        } catch (e) {
+          // A failed step is transactional (see generate.ts) — safe to
+          // re-forward the same token on the graph path below.
+          this.#compiled = null;
+          console.warn(`batch lane: compiled decode disabled: ${e}`);
+        } finally {
+          if (owned) cur.dispose();
+        }
+      }
+      let fwd: Cache[] | null = null;
       try {
-        const ids = this.#pendingToks
-          ? ops.reshape(this.#pendingToks, [B, 1]) // feed the unread tokens
-          : ops.fromInt32(rows.map((r) => r.current), [B, 1]); // pipeline cold
-        const h = this.model.forwardHidden(ids, fwd);
-        ids.dispose();
-        const lg = this.model.logitsFromHidden(h); // [B,1,V]
-        h.dispose();
+        if (!lg) {
+          fwd = inners.map((c) =>
+            c instanceof BatchedRotatingCache || c instanceof BatchedRotatingQuantCache ||
+            c instanceof SSMCache || unpadded
+              ? c
+              : c instanceof QuantizedKVCache
+                ? new BatchedQuantDecodeMaskCache(c, B, this.#fullLeftPad)
+                : new BatchedDecodeMaskCache(c, B, this.#fullLeftPad, null),
+          );
+          const ids = this.#pendingToks
+            ? ops.reshape(this.#pendingToks, [B, 1]) // feed the unread tokens
+            : ops.fromInt32(rows.map((r) => r.current), [B, 1]); // pipeline cold
+          const h = this.model.forwardHidden(ids, fwd);
+          ids.dispose();
+          lg = this.model.logitsFromHidden(h); // [B,1,V]
+          h.dispose();
+        }
         const V = lg.shape[2]!;
-        const sampled: MlxArray[] = [];
-        for (let b = 0; b < B; b++) {
-          const row = rows[b]!;
-          if (row.sampled >= row.req.maxTokens) {
-            // Length-doomed row: evicted right after the emission below ever
-            // uses this slot as input — placeholder keeps the [B] alignment.
-            sampled.push(ops.fromInt32([0], [1]));
-            continue;
+        // Vectorized homogeneous sampling (batching-perf-path P0): when every
+        // LIVE row is plain greedy, one log-softmax+argmax over [B,V] replaces
+        // B slice/sample/concat graphs. Per-row identical math (row-independent
+        // ops, same per-row shapes — argmax over log-softmax mirrors the
+        // closure's toLogprobs→argmax exactly, tie behavior included). Doomed
+        // rows get a real argmax instead of the placeholder 0 — equally
+        // harmless (the slot is filtered before it is ever emitted; its only
+        // use is one KV write on the row's own, about-to-evict row).
+        const vecOk =
+          process.env.MLX_BUN_BATCH_VEC_SAMPLE !== "0" &&
+          rows.every((r) => r.sampled >= r.req.maxTokens || r.req.plainGreedy);
+        // Doomed slots hold placeholders (vec path: a real argmax value,
+        // equally not part of the row's stream) — flag them so the fed
+        // accounting at the read stays per-row exact.
+        nextReal = rows.map((r) => r.sampled < r.req.maxTokens);
+        if (vecOk) {
+          const flat = ops.reshape(lg, [B, V]);
+          lg.dispose();
+          const lp = toLogprobs(flat);
+          flat.dispose();
+          nextToks = ops.argmaxAxis(lp, -1); // [B]
+          lp.dispose();
+          for (const row of rows) if (row.sampled < row.req.maxTokens) row.sampled++;
+        } else {
+          const sampled: MlxArray[] = [];
+          for (let b = 0; b < B; b++) {
+            const row = rows[b]!;
+            if (row.sampled >= row.req.maxTokens) {
+              // Length-doomed row: evicted right after the emission below ever
+              // uses this slot as input — placeholder keeps the [B] alignment.
+              sampled.push(ops.fromInt32([0], [1]));
+              continue;
+            }
+            const rl = lg.slice([b, 0, 0], [b + 1, 1, V]);
+            const rl2 = ops.reshape(rl, [1, V]);
+            rl.dispose();
+            sampled.push(row.req.sample(rl2, row.sampled));
+            row.sampled++;
+            rl2.dispose();
           }
-          const rl = lg.slice([b, 0, 0], [b + 1, 1, V]);
-          const rl2 = ops.reshape(rl, [1, V]);
-          rl.dispose();
-          sampled.push(row.req.sample(rl2, row.sampled));
-          row.sampled++;
-          rl2.dispose();
+          lg.dispose();
+          nextToks = ops.concatAxis(sampled, 0); // [B]
+          for (const t of sampled) t.dispose();
         }
-        lg.dispose();
-        nextToks = ops.concatAxis(sampled, 0); // [B]
-        for (const t of sampled) t.dispose();
-        ops.asyncEvalAll([nextToks]); // dispatch; read NEXT iteration
+        // dispatch; read NEXT iteration. evalWith: the compiled step's
+        // cache-update nodes must ride the same async_eval (generate.ts).
+        ops.asyncEvalAll([nextToks, ...evalWith]);
       } finally {
         // Free the step's RoPE arrays; do NOT dispose (full wrappers would free
         // their persistent inner; rot caches persist across steps).
-        for (const c of fwd) (c as { releaseRopeArr?: () => void }).releaseRopeArr?.();
+        if (fwd) for (const c of fwd) (c as { releaseRopeArr?: () => void }).releaseRopeArr?.();
       }
     }
     this.#steps++;
     if (this.#steps % 256 === 0) clearCache(); // serial's cadence, not per-step
+    if (STEP_TRACE) STEP_T.build += performance.now() - STEP_T.t0;
 
     // Read + emit the PREVIOUS step's tokens while the new step computes.
     const prev = this.#pendingToks;
+    const prevReal = this.#pendingReal;
     this.#pendingToks = nextToks;
+    this.#pendingReal = nextReal;
     // Kill switch / A-B lever (house style, cf. MLX_BUN_COMPILED_DECODE=0):
     // MLX_BUN_BATCH_NO_PIPELINE=1 reads THIS step's tokens synchronously —
     // set from process start `prev` is always null, so the flush below IS the
@@ -446,10 +1091,23 @@ export class BatchScheduler {
       return;
     }
     if (prev) {
-      const toks = [...prev.toFloat32()].map((x) => Math.round(x));
+      const tRead = STEP_TRACE ? performance.now() : 0;
+      const toks = prev.toIntTokens();
       prev.dispose();
+      // These values were the step's forward input (fed) iff a forward ran.
+      // Placeholder slots are excluded from `fed` — and taint their row if
+      // they ever fed (see Row.fedTainted; structurally unreachable today).
+      if (anyLive)
+        for (let b = 0; b < B; b++) {
+          if (!prevReal || prevReal[b]) rows[b]!.fed.push(toks[b]!);
+          else rows[b]!.fedTainted = true;
+        }
+      if (STEP_TRACE) STEP_T.read += performance.now() - tRead;
+      const tEmit = STEP_TRACE ? performance.now() : 0;
       await this.#emitRows(toks); // also filters #pendingToks on eviction
+      if (STEP_TRACE) { STEP_T.emit += performance.now() - tEmit; STEP_T.n++; }
     }
+    if (STEP_TRACE) STEP_T.lastEnd = performance.now();
   }
 
   /** Read out the pipeline register (if any): emit its tokens and evict
@@ -458,17 +1116,149 @@ export class BatchScheduler {
     const prev = this.#pendingToks;
     if (!prev) return;
     this.#pendingToks = null;
-    const toks = [...prev.toFloat32()].map((x) => Math.round(x));
+    this.#pendingReal = null; // flushed values never fed — nothing to account
+    const toks = prev.toIntTokens();
     prev.dispose();
+    // Grammar rows: the flushed tokens are EMITTED below, so their matchers
+    // must advance here — #stepGrammar's accept only covers tokens it reads
+    // from a live pending array, and after this flush it cold-starts (no
+    // accept). Skipping this left the matcher one token behind its stream on
+    // every mid-decode join → one-step-stale masks → invalid output (found
+    // by the feature-matrix conformance gate, 2026-07-03; the fill fired
+    // here is awaited by the next #stepGrammar's ready()).
+    for (let b = 0; b < this.#running.length && b < toks.length; b++) {
+      const g = this.#running[b]!.req.grammar;
+      if (g && !g.isTerminated) g.accept(toks[b]!);
+    }
     await this.#emitRows(toks);
   }
 
+  /** The read-before-build decode step for batches with ≥1 live grammar row
+   *  (B1). Mirrors the serial loop's grammar resolution: read the previous
+   *  step's [B] tokens NOW (acceptToken needs JS numbers), accept() per row
+   *  (fires async fills), build the forward graph (fills overlap), await
+   *  ready(), sample per live row (the closure applies applyMask), then emit
+   *  the values read in step 1 — no second readback. Terminated grammar rows
+   *  keep their [B] slot through the forward (placeholder, not sampled) and
+   *  finish("stop") in #emitRows. On a cold start (first step after prefill)
+   *  there is no pending array to read; the prefill already accepted token 0
+   *  and fired the fill, so we just await ready() + sample. */
+  async #stepGrammar(): Promise<void> {
+    const rows = this.#running;
+    const B = rows.length;
+    const inners = this.#inners!;
+
+    // (1) Read the pending [B] token array (host copy for accept). The device
+    //     array is kept for the forward input below; disposed after sampling.
+    //     Cold start: pendingToks null → prevVals empty (prefill handled tok0).
+    const prev = this.#pendingToks;
+    const prevVals: number[] =
+      prev ? prev.toIntTokens() : [];
+    if (process.env.MLX_BUN_GRAMMAR_DEBUG === "1")
+      console.log(`[sg] B=${B} prevVals=${JSON.stringify(prevVals)} current=${JSON.stringify(rows.map((r) => r.current))} sampled=${JSON.stringify(rows.map((r) => r.sampled))}`);
+
+    // (2) accept() per live grammar row — fires that row's async bitmask fill.
+    for (let b = 0; b < B && prevVals.length; b++) {
+      const g = rows[b]!.req.grammar;
+      if (g && !g.isTerminated) g.accept(prevVals[b]!);
+    }
+
+    let nextToks: MlxArray | null = null;
+    let nextReal: boolean[] | null = null;
+    const anyLive = rows.some((r) => r.sampled < r.req.maxTokens);
+    if (anyLive) {
+      // Fed-token accounting (prompt-cache put) — mirror of #step's read:
+      // real slots push, a placeholder that fed taints its row.
+      const prevReal = this.#pendingReal;
+      if (prev)
+        for (let b = 0; b < B; b++) {
+          if (!prevReal || prevReal[b]) rows[b]!.fed.push(prevVals[b]!);
+          else rows[b]!.fedTainted = true;
+        }
+      else for (const r of rows) r.fed.push(r.current);
+      // Unpadded fast path — same rule as #step (bare caches == serial graph).
+      const unpadded = this.#fullLeftPad.every((p) => p === 0);
+      const fwd: Cache[] = inners.map((c) =>
+        c instanceof BatchedRotatingCache || c instanceof BatchedRotatingQuantCache ||
+            c instanceof SSMCache || unpadded
+          ? c
+          : c instanceof QuantizedKVCache
+            ? new BatchedQuantDecodeMaskCache(c, B, this.#fullLeftPad)
+            : new BatchedDecodeMaskCache(c, B, this.#fullLeftPad, null),
+      );
+      try {
+        // (3) Build the forward graph (host-side; the fills overlap it).
+        const ids = prev
+          ? ops.reshape(prev, [B, 1]) // feed the unread tokens (device array)
+          : ops.fromInt32(rows.map((r) => r.current), [B, 1]); // pipeline cold
+        const h = this.model.forwardHidden(ids, fwd);
+        ids.dispose();
+        const lg = this.model.logitsFromHidden(h); // [B,1,V]
+        h.dispose();
+        const V = lg.shape[2]!;
+
+        // (4) await ready() on every live grammar row — the fills fired in (2)
+        //     overlapped the graph build above. Usually already resolved.
+        const liveGrammar = rows.filter(
+          (r) => r.req.grammar && !r.req.grammar.isTerminated,
+        );
+        if (liveGrammar.length)
+          await Promise.all(liveGrammar.map((r) => r.req.grammar!.ready()));
+
+        // (5) Sample per live row (the closure applies applyMask after the
+        //     logits processors). Terminated grammar rows + length-doomed rows
+        //     take a placeholder (one harmless KV write, then evicted).
+        const sampled: MlxArray[] = [];
+        nextReal = rows.map(
+          (r) => r.sampled < r.req.maxTokens && !r.req.grammar?.isTerminated,
+        );
+        for (let b = 0; b < B; b++) {
+          const row = rows[b]!;
+          if (
+            row.sampled >= row.req.maxTokens ||
+            row.req.grammar?.isTerminated
+          ) {
+            sampled.push(ops.fromInt32([0], [1]));
+            continue;
+          }
+          const rl = lg.slice([b, 0, 0], [b + 1, 1, V]);
+          const rl2 = ops.reshape(rl, [1, V]);
+          rl.dispose();
+          sampled.push(row.req.sample(rl2, row.sampled));
+          row.sampled++;
+          rl2.dispose();
+        }
+        lg.dispose();
+        nextToks = ops.concatAxis(sampled, 0); // [B]
+        for (const t of sampled) t.dispose();
+        ops.asyncEvalAll([nextToks]);
+      } finally {
+        for (const c of fwd)
+          (c as { releaseRopeArr?: () => void }).releaseRopeArr?.();
+      }
+    }
+    this.#steps++;
+    if (this.#steps % 256 === 0) clearCache();
+
+    // The device array `prev` has now fed the forward + been read for accept;
+    // dispose it and install the new pending array. On cold start prev is null.
+    prev?.dispose();
+    this.#pendingToks = nextToks;
+    this.#pendingReal = nextReal;
+
+    // (6) Emit the values read in (1) — no second readback. Terminated grammar
+    //     rows finish("stop") via #emit's isTerminated check; the filter evicts.
+    //     Cold start: prevVals empty (prefill emitted tok0) → nothing to emit.
+    if (prevVals.length) await this.#emitRows(prevVals);
+  }
+
   /** Emit one read-back token per running row; evict finished rows. A row's
    *  onToken throwing rejects THAT row and evicts it — siblings continue. */
   async #emitRows(toks: number[]): Promise<void> {
     const rows = this.#running;
     const B = rows.length;
     const keep: number[] = [];
+    const done: { b: number; row: Row }[] = []; // clean finishes (stop/length)
     for (let b = 0; b < B; b++) {
       const row = rows[b]!;
       row.generated++;
@@ -477,12 +1267,27 @@ export class BatchScheduler {
         disp = await this.#emit(row, toks[b]!);
       } catch (e) {
         row.reject(e); // containment: this row only (mlx-lm `remove`)
+        row.merged = true; // poison: a rejected row is never put() back
+        row.fedTainted = true;
         continue; // not kept → evicted by the filter below
       }
       if (disp === "continue") keep.push(b);
-      else this.#finish(row, disp);
+      else {
+        done.push({ b, row });
+        this.#finish(row, disp);
+      }
+    }
+    if (keep.length < B) {
+      // Per-row prompt-cache extraction (mlx-lm server.py:864-880 →
+      // BatchGenerator.extract_cache): a MERGED row's KV lives interleaved
+      // in the batched inners, so it is pulled out into fresh serial caches
+      // BEFORE filter() mutates them. Never-merged lone rows skip this —
+      // #applyFilter's keep=[] path put()s their adopted caches zero-copy.
+      // Rejected rows aren't in `done`; the whole-batch error path calls
+      // #applyFilter(dropOnly) directly and never reaches here.
+      for (const { b, row } of done) this.#extractAndPut(b, row);
+      this.#applyFilter(keep);
     }
-    if (keep.length < B) this.#applyFilter(keep);
   }
 
   /** Account one sampled token for a row. Mirrors generate(): EOS terminates
@@ -492,6 +1297,10 @@ export class BatchScheduler {
     if (row.req.eosTokenIds.includes(token)) return "stop";
     const cont = await row.req.onToken(token);
     if (cont === false) return "stop";
+    // Grammar termination: the matcher is satisfied (e.g. closing `}` accepted).
+    // The final token has been delivered via onToken above; halt with "stop" so
+    // the row finishes + evicts rather than sampling into an all--inf mask.
+    if (row.req.grammar?.isTerminated) return "stop";
     if (row.generated >= row.req.maxTokens) return "length";
     row.current = token;
     return "continue";
@@ -501,22 +1310,119 @@ export class BatchScheduler {
     row.resolve({
       promptTokens: row.promptTokens,
       generatedTokens: row.generated,
-      cachedTokens: 0,
+      cachedTokens: row.cachedTokens,
       finishReason: reason,
     });
   }
 
+  /** Finish-time disposition of serial-class caches covering exactly
+   *  `tokens`: put() into the prompt cache when the hook is present and the
+   *  offset lines up (defensive — a mismatch means the accounting is wrong
+   *  and the entry would corrupt future hits), else dispose. `retain` rides
+   *  along per the PromptCacheEntry contract (runs after dispose). */
+  #putOrDispose(caches: Cache[], tokens: number[], retain?: () => void): void {
+    if (this.#promptCache) {
+      const withOff = caches.find(
+        (c) => typeof (c as { offset?: unknown }).offset === "number",
+      ) as { offset: number } | undefined;
+      if (withOff && withOff.offset === tokens.length) {
+        this.#promptCache.put(tokens, caches, "", retain);
+        return;
+      }
+    }
+    for (const c of caches) c.dispose();
+    retain?.();
+  }
+
+  /** Extract a finishing MERGED row's KV into fresh serial caches and put()
+   *  them keyed by [promptIds + fed] — the batch-lane mirror of mlx-lm
+   *  server.py:872 (extract_cache → prompt_cache.insert_cache). Gates:
+   *  promptTokens >= 256 (the boundary-snapshot substantiality gate,
+   *  server.ts) and an exact coverage key (!fedTainted). Refusal
+   *  (#extractRowCaches null) just disposes-by-omission — the row's KV dies
+   *  with the filter, exactly the pre-extraction behavior. */
+  #extractAndPut(b: number, row: Row): void {
+    if (!this.#promptCache || !row.merged || row.fedTainted) return;
+    if (row.promptTokens < 256) return;
+    const tokens = [...row.req.promptIds, ...row.fed];
+    const caches = this.#extractRowCaches(b, tokens.length);
+    if (!caches) return;
+    // Materialize the owned copies without stalling the in-flight step (the
+    // slices depend on this step's KV writes): async — the batched source
+    // buffers free once the copies land, instead of being pinned by a lazy
+    // graph inside an idle cache entry.
+    ops.asyncEvalAll(caches.flatMap((c) => c.state()));
+    this.#putOrDispose(caches, tokens);
+  }
+
+  /** Row `b` of every layer as OWNED serial-class caches, or null when a
+   *  layer kind can't be extracted (then the caller drops the row's KV as
+   *  before). Bit-exactness: merge/extend/filter/decode are byte-preserving
+   *  per row (tests/batched-decode-parity, tests/batched-rotating,
+   *  tests/batched-rotating-quant) and each extract is a pure slice+copy of
+   *  those bytes (tests/batched-extract), so an extracted row's bytes ==
+   *  the solo run's. */
+  #extractRowCaches(b: number, expectTokens: number): Cache[] | null {
+    const out: Cache[] = [];
+    for (const inner of this.#inners!) {
+      let c: Cache | null;
+      if (inner instanceof BatchedRotatingQuantCache) c = inner.extractRow(b);
+      else if (inner instanceof BatchedRotatingCache) c = inner.extractRow(b);
+      else if (inner instanceof SSMCache)
+        // Coverage gate: recurrent state is UNTRIMMABLE, so an entry is only
+        // valid when the row's own advance count equals the [promptIds+fed]
+        // key EXACTLY (a mismatched entry would silently corrupt every future
+        // exact-hit). Defensive — a miss here means the fed accounting broke.
+        c = inner.conv && inner.rowOffset(b) === expectTokens ? inner.extractRow(b) : null;
+      else if (
+        inner instanceof RotatingQuantizedKVCache || // adopted serial state never
+        inner instanceof RotatingKVCache //     coexists with a merged row (defensive)
+      )
+        c = null;
+      else if (inner instanceof QuantizedKVCache)
+        c = inner.keys ? extractQuantRow(inner, this.#fullLeftPad[b]!, b) : null;
+      else c = inner.keys ? extractKVRow(inner, this.#fullLeftPad[b]!, b) : null;
+      if (!c) {
+        for (const d of out) d.dispose();
+        return null;
+      }
+      out.push(c);
+    }
+    return out;
+  }
+
   /** Evict rows not in `keep` (sorted ascending) from the batched KV and the
    *  pipeline register. */
-  #applyFilter(keep: number[]): void {
+  #applyFilter(keep: number[], dropOnly = false): void {
     const inners = this.#inners!;
     if (keep.length === 0) {
-      for (const c of inners) c.dispose();
+      // Prompt-cache put (Phase 3.2): a lone NEVER-MERGED row's inners are
+      // its adopted serial-class caches, covering exactly prompt+fed — hand
+      // them back to the cache instead of disposing. dropOnly (the batch-
+      // drop error path) and poisoned rows dispose as before; MERGED rows'
+      // KV was already extracted per row in #emitRows (owned copies), so
+      // disposing the batched inners here is safe either way.
+      const solo =
+        !dropOnly && this.#running.length === 1 && !this.#running[0]!.merged
+          ? this.#running[0]!
+          : null;
+      if (solo) {
+        this.#putOrDispose(
+          inners as Cache[],
+          [...solo.req.promptIds, ...solo.fed],
+          this.#adoptedRetain ?? undefined,
+        );
+      } else {
+        for (const c of inners) c.dispose();
+        this.#adoptedRetain?.();
+      }
+      this.#adoptedRetain = null;
       this.#inners = null;
       this.#fullLeftPad = [];
       this.#running = [];
       this.#pendingToks?.dispose();
       this.#pendingToks = null;
+      this.#pendingReal = null;
       return;
     }
     const out: LayerInner[] = [];
@@ -524,6 +1430,24 @@ export class BatchScheduler {
       if (inner instanceof BatchedRotatingCache) {
         inner.filter(keep); // in-place (mutates + drops rows + reduces padding)
         out.push(inner);
+      } else if (inner instanceof BatchedRotatingQuantCache) {
+        inner.filter(keep); // in-place B-axis take on all six components
+        out.push(inner);
+      } else if (inner instanceof RotatingQuantizedKVCache || inner instanceof RotatingKVCache) {
+        // Adopted lone-row state exists only at B=1, where the only filter
+        // is the keep=[] dispose-all handled above — unreachable.
+        throw new Error("applyFilter: adopted serial rotating cache cannot be row-filtered");
+      } else if (inner instanceof SSMCache) {
+        inner.filter(keep); // in-place B-axis take on both state slots
+        out.push(inner);
+      } else if (inner instanceof QuantizedKVCache) {
+        const [k0, v0] = inner.temporalView();
+        const f = filterQuantRows(k0, v0, keep);
+        for (const t of [k0, v0]) { t.packed.dispose(); t.scales.dispose(); t.biases.dispose(); }
+        const c = new QuantizedKVCache(inner.groupSize, inner.bits);
+        c.restoreState(f.keys, f.values, inner.offset);
+        out.push(c);
+        inner.dispose();
       } else {
         const [k0, v0] = inner.temporalView();
         const f = filterKVRows(k0, v0, keep);
@@ -544,10 +1468,11 @@ export class BatchScheduler {
       this.#pendingToks.dispose();
       this.#pendingToks = next;
     }
+    if (this.#pendingReal) this.#pendingReal = keep.map((i) => this.#pendingReal![i]!);
   }
 
   #readToken(t: MlxArray): number {
-    const v = Math.round(t.toFloat32()[0]!);
+    const v = t.toIntTokens()[0]!;
     t.dispose();
     return v;
   }